I think the optimizers and schedulers could use some extra work.
The main thing I would love, is for optimizers and schedulers to have a to() method so we can send their parameters to a certain device. This would allow us to save the parameters and reload them at a later stage without any problems.

Right now, the following script crashes with an error that the optimizer expects a torch.FloatTensor but got a torch.cuda.FloatTensor. I am not quite sure why, as I specifically load the data to CPU?

# Note: This is not my actual script, but rather a representation of the steps and in what order I perform them
network = ... # Creating a network based of torch.nn.Module
optim = torch.optim.SGD(network.parameters(), ...)

state = torch.load('file_path', lambda storage, loc: storage)  # Send all tensors to CPU

network.load_state_dict(state['network'])
optim.load_state_dict(state['optim'])

network.to(torch.device('cuda'))
#optim.to(torch.device('cuda'))   # Would this solve my problems? 

The same could be said about schedulers.