Load models

[chammami]

Thanks for this great project.

I started using Tez con Kaggle competition and found that model loading is not appropriately handled when training/infering on different devices: GPU -> CPU. This one possible solution

    def load(self, model_path, device="cuda"):
        self.device = device
        if next(self.parameters()).device != self.device:
            self.to(self.device)
        model_dict = torch.load(model_path, map_location=torch.device(device))
        self.load_state_dict(model_dict["state_dict"])

[abhishekkrthakur]

should be fixed now