(when checking argument for argument mat1 in method wrapper_addmm)

Jun_Park: <p>class NICE(pl.LightningModule):</p>
<pre><code>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))
    # logistic distribution for prior
    base = torch.distributions.uniform.Uniform(torch.tensor(0.0).to('cuda'), torch.tensor(1.0).to('cuda'))
    transforms = [torch.distributions.transforms.SigmoidTransform().inv,
                  torch.distributions.transforms.AffineTransform(torch.tensor(0.0).to('cuda'), 
                                                                 torch.tensor(1.0).to('cuda'))]
    self.prior = torch.distributions.TransformedDistribution(base, transforms)

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
</code></pre>
<p>========================<br>
It works for training, but error occurs when sampling:</p>
<p>import matplotlib.pyplot as plt</p>
<p>samples = model.sample(img_shape=[16,1,28,28])<br>
samples = samples.view(16, 28, 28).cpu().numpy()<br>
plt.figure(figsize=(10,10))<br>
for i in range(16):<br>
plt.subplot(4,4,i+1)<br>
plt.imshow(samples[i], cmap=“gray”)<br>
plt.axis(“off”)<br>
plt.show()</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/4/b4012d60cfb38e57b4fe3e82d296d5e883cc2caa.png" data-download-href="https://discuss.pytorch.org/uploads/default/b4012d60cfb38e57b4fe3e82d296d5e883cc2caa" title="Screen Shot 2023-01-16 at 4.31.44 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/4/b4012d60cfb38e57b4fe3e82d296d5e883cc2caa_2_690x482.png" alt="Screen Shot 2023-01-16 at 4.31.44 PM" data-base62-sha1="pGosWWbg8wui2LJOl4dsDjXM5Bg" width="690" height="482" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/4/b4012d60cfb38e57b4fe3e82d296d5e883cc2caa_2_690x482.png, https://discuss.pytorch.org/uploads/default/original/3X/b/4/b4012d60cfb38e57b4fe3e82d296d5e883cc2caa.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/b/4/b4012d60cfb38e57b4fe3e82d296d5e883cc2caa.png 2x" data-dominant-color="F2F5F4"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screen Shot 2023-01-16 at 4.31.44 PM</span><span class="informations">743×520 73.9 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

ptrblck: <p>Check if <code>self.layers</code> were moved to the GPU and if <code>x2</code> was moved as well since one of these two objects is still stored on the CPU and causes the device mismatch.</p>

Jun_Park: <p>Thank you for the prompt reply. I checked that x1 and x2 are on the GPU and moved self.layers to the GPU using the following code. <strong>Still get the same error.</strong></p>
<p>def <strong>init</strong>(self, in_features=784, hidden_features=1000, num_coupling=4):<br>
super().<strong>init</strong>()<br>
self.save_hyperparameters()<br>
self.layers = nn.ModuleList()<br>
self.num_coupling = num_coupling<br>
for _ in range(self.num_coupling):<br>
self.layers.append(nn.Sequential(<br>
nn.Linear(in_features // 2, hidden_features), nn.ReLU(),<br>
nn.Linear(hidden_features, hidden_features), nn.ReLU(),<br>
nn.Linear(hidden_features, hidden_features), nn.ReLU(),<br>
nn.Linear(hidden_features, hidden_features), nn.ReLU(),<br>
nn.Linear(hidden_features, hidden_features), nn.ReLU(),<br>
nn.Linear(hidden_features, in_features // 2),))<br>
<strong>self.layers.to(‘cuda’)</strong><br>
self.scale = nn.Parameter(torch.zeros(in_features)).to(‘cuda’)</p>
<p>BTW, I don’t understand why this code works in training but not in sampling. I thought that self.layers in the model moved to the GPU when I initialized the model by</p>
<p>model = NICE().to(‘cuda’)</p>
<p>After the training, isn’t the model still staying on the GPU?</p>

ptrblck: <p>Yes, you are right. Once the model is moved to the GPU its submodules and parameters will stay there unless you are reinitializing it. Since your training seems to work whole the inference fails could you check why the model changed? Are you recreating it somewhere?</p>

Jun_Park: <p>Not that I know of. Sampling works now by adding the self.layers.to(‘cuda’) in the inverse_ function as you’ve advised.</p>
<p>def inverse(self, z):<br>
# z: (batch_size, 784)<br>
z = z.view(z.size(0), -1)<br>
x = self.inverse_(z)<br>
x = x.view(x.size(0), 1, 28, 28)<br>
return x</p>
<pre><code>def inverse_(self, z):
    # z: (batch_size, 784)
    x = z
    s = torch.exp(self.scale).to('cuda')
    x = x / s
    **self.layers.to('cuda')** 
    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
</code></pre>
<p>These three functions (methods) were not used in the training step since only the forward step was used. I guess the sampling call might recreate the model.</p>
<p>samples = model.sample(img_shape=[16,1,28,28])</p>
<p>But still, I don’t get the picture.</p>

ptrblck: <p>I also cannot explain the issue you are seeing since your code works fine for me:</p>
<pre><code class="lang-python">model = NICE().to('cuda')

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

samples = model.sample(img_shape=[16,1,28,28])
print(samples.device)
# cuda:0
</code></pre>
<p>Note that I had to remove the <code>pl.LightningModule</code> as the base class and used the plain <code>nn.Module</code> instead. Do you see the error using my code snippet and if so, could you also replace the <code>LightningModule</code> with <code>nn.Module</code>?</p>

Jun_Park: <p>I think the culprit is pl.Trainer which automatically assigns a device for training. I ran your code before and after the training via trainer.fit and found the following,</p>
<h1>
<a name="before-training-the-model-is-on-gpu-1" class="anchor" href="#before-training-the-model-is-on-gpu-1"></a>Before training; the model is on GPU.</h1>
<p>x = torch.randn(10, 784).to(‘cuda’)<br>
out = model(x)<br>
print(out.device) # cuda:0</p>
<h1>
<a name="init-data-2" class="anchor" href="#init-data-2"></a>Init data</h1>
<p>dm = MNISTDataModule(batch_size=64)</p>
<h1>
<a name="init-trainer-3" class="anchor" href="#init-trainer-3"></a>Init trainer</h1>
<p>trainer = pl.Trainer(default_root_dir=os.path.join(CHECKPOINT_PATH, “NICE”),<br>
accelerator=“gpu”,<br>
devices=1,<br>
max_epochs=1,<br>
gradient_clip_val=1.0,<br>
callbacks=[ModelCheckpoint(save_weights_only=True, mode=“min”, monitor=“val_loss”),<br>
TQDMProgressBar(refresh_rate=50)])</p>
<h1>
<a name="train-4" class="anchor" href="#train-4"></a>Train</h1>
<p>trainer.fit(model, dm)</p>
<h1>
<a name="after-training-model-is-on-cpu-5" class="anchor" href="#after-training-model-is-on-cpu-5"></a>After training; model is on CPU.</h1>
<p>x = torch.randn(10, 784)<br>
out = model(x)<br>
print(out.device) # cpu</p>
<p>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.</p>
<p>Thank you so much for the help!</p>

ptrblck: <aside class="quote no-group" data-username="Jun_Park" data-post="7" data-topic="170352">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/jun_park/48/56445_2.png" class="avatar"> Jun_Park:</div>
<blockquote>
<p>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.</p>
</blockquote>
</aside>
<p>That’s interesting. Do you know why that’s the case as it seems to be a bit unexpected (at least to me)?</p>

Jun_Park: <p>I have no idea^^. I think it is related to the fact that Pytorch Lightning takes care of device assignments by itself. No further clue.</p>
<ol start="4">
<li>There are no <code>.cuda()</code> or <code>.to(device)</code> calls required. Lightning does these for you.</li>
</ol>
<h1>
<a name="dont-do-in-lightning-1" class="anchor" href="#dont-do-in-lightning-1"></a>don’t do in Lightning</h1>
<p>x = torch.Tensor(2, 3)<br>
x = x.cuda()<br>
x = x.to(device)</p>
<h1>
<a name="do-this-instead-2" class="anchor" href="#do-this-instead-2"></a>do this instead</h1>
<p>x = x</p>
<h1>
<a name="leave-it-alone-or-to-init-a-new-tensor-3" class="anchor" href="#leave-it-alone-or-to-init-a-new-tensor-3"></a>leave it alone! # or to init a new tensor</h1>
<p>new_x = torch.Tensor(2, 3)<br>
new_x = new_x.to(x)</p>