Thanks @albanD for replying! Yes, I called self.model.cuda() when the model is initialized. I thought that I can switch it back to CPU by calling self.network.cpu(). Can’t I?
The reason why I’m doing this (switching GPU <> CPU) is because I was getting GPU out of memory error and start, end, _ = self.network(batch) which is a simple forward propagation (all autograd disabled by calling self.network.eval().)