Cubbee:

pred = torch.from_numpy(self.model(*input_batch).cpu().data.numpy())


I have the same problem here. When I do .cpu() for network output, by each iteration gpu memory starts to fill up. Did you find a more normative solution ?