Where are the tensors (model) after training via pl.Trainer?

[junpark25]

I am asking because I've got the following error when I ran the following code after training.

samples = model.sample(img_shape=[16,1,28,28]) 

I checked that before training the model was on the GPU.

x = torch.randn(10, 784).to(‘cuda’)
out = model(x)
print(out.device) # cuda:0

But after training

# Init data
dm = MNISTDataModule(batch_size=64)

# Init trainer
trainer = pl.Trainer(default_root_dir=os.path.join(CHECKPOINT_PATH, “NICE”),
  accelerator=“gpu”,
  devices=1,
  max_epochs=1,
  gradient_clip_val=1.0,
  callbacks=[ModelCheckpoint(save_weights_only=True, mode=“min”, monitor=“val_loss”),
  TQDMProgressBar(refresh_rate=50)])

Train
trainer.fit(model, dm)

I found that model was on the CPU.

x = torch.randn(10, 784)
out = model(x)
print(out.device) # cpu

It seems that pl.Trainer moves the tensors (model) to the GPU during training and back to the CPU after training, if their locations are not explicitly specified. And this caused the error.

Is this normal behavior or I messed with some configuration?

My lightning model is

class NICE(pl.LightningModule):
    def __init__(self, in_features=784, hidden_features=1000, num_coupling=4):
        super().__init__()
        self.save_hyperparameters()
        self.layers = nn.ModuleList()
        self.num_coupling = num_coupling
        for _ in range(self.num_coupling):
            self.layers.append(nn.Sequential(
                nn.Linear(in_features // 2, hidden_features), nn.ReLU(),
                nn.Linear(hidden_features, hidden_features), nn.ReLU(),
                nn.Linear(hidden_features, hidden_features), nn.ReLU(),
                nn.Linear(hidden_features, hidden_features), nn.ReLU(),
                nn.Linear(hidden_features, hidden_features), nn.ReLU(),
                nn.Linear(hidden_features, in_features // 2),))
        self.scale = nn.Parameter(torch.zeros(in_features))
        self.prior = torch.distributions.normal.Normal(loc=0.0, scale=1.0) # Gaussian prior
    
    def forward(self, x):
        # x: (batch_size, 784)
        x = x.view(x.size(0), -1)
        z, log_det = self.forward_(x)
        log_prob = self.prior.log_prob(z).sum(dim=1) + log_det
        return log_prob
    
    def forward_(self, x):
        # x: (batch_size, 784)
        z = x
        log_det = 0
        s = torch.exp(self.scale) # Positive scale
        for i in range(self.num_coupling):
            # Check the number of layers in self.layers
            # print(len(self.layers)) # 4
            z1 = z.chunk(2, dim=1)[0] if (i % 2 == 0) else z.chunk(2, dim=1)[1]
            z2 = z.chunk(2, dim=1)[1] if (i % 2 == 0) else z.chunk(2, dim=1)[0]
            z2 = z2 + self.layers[i](z1)
            z = torch.cat([z1, z2], dim=1) if (i % 2 == 0) else torch.cat([z2, z1], dim=1)
        z = z * s
        log_det += torch.log(s).sum()
        return z, log_det
    
    def inverse(self, z):
        # z: (batch_size, 784)
        z = z.view(z.size(0), -1)
        x = self.inverse_(z)
        x = x.view(x.size(0), 1, 28, 28)
        return x
    
    def inverse_(self, z):
        # z: (batch_size, 784)
        x = z
        s = torch.exp(self.scale).to('cuda')
        x = x / s
        # print(x.is_cuda) # True
        for i in range(self.num_coupling -1, -1, -1):
            # Assumed that the num_coupling is even
            x1 = x.chunk(2, dim=1)[0] if (i % 2 == 1) else x.chunk(2, dim=1)[1]
            x2 = x.chunk(2, dim=1)[1] if (i % 2 == 1) else x.chunk(2, dim=1)[0]
            x1 = x1 - self.layers[i](x2)
            x = torch.cat([x1, x2], dim=1) if (i % 2 == 1) else torch.cat([x2, x1], dim=1)
        return x
    
    def sample(self, img_shape):
        z = self.prior.sample(img_shape).to('cuda')
        x = self.inverse(z)
        return x