I am having trouble figuring out what to put on the device and how.
I noticed that everything I declare in the constructor goes to device when I call model.to(device).
def init(self, latent_vector_size):
self.latent_vector_size = latent_vector_size
self.encoder = nn.Sequential(
self.laten_linear = nn.Linear(3136, 2)
In above both encoder and laten_linear went to device.
But when I define this method, it doesn’t go to device:
def sample_z(self, mean, logvar):
stddev = torch.exp(0.5 * logvar)
noise = torch.randn(stddev.size())
return (noise * stddev) + mean
How do I put this method on the device?
Also is there any rule of thumb as to what goes on the device?