Automatic mixed precision - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument mat1 in method wrapper_addmm) error

username: Nimrod_Daniel
cooked: <p>I try to do automatic mixed-precision, but unfortunately I get an error that says that all tensors are not on the GPU. When I check the model and all variables it seems ok.<br>
Any idea why it fails during the forward pass, and how to fix it?<br>
(I chopped a lot of code so it’d be the shortest to reproduce the error)<br>
(torch 1.10.2, python 3.6.9, T4 GPU)</p>
<p>import torch<br>
print(torch.<strong>version</strong>)</p>
<p>import os<br>
import torch.nn as nn<br>
from torch.optim import SGD<br>
from torchvision import transforms,models<br>
from torchvision.datasets import ImageFolder<br>
from torch.cuda.amp import GradScaler, autocast</p>
<p>dataset_path = ‘/my_source_path’ # under the path I have ‘train’ and ‘test’ folders</p>
<p>device = torch.device(“cuda:0” if torch.cuda.is_available() else “cpu”)<br>
print('device: ', device)</p>
<p>train_transform = transforms.Compose([<br>
transforms.Resize(224),<br>
transforms.ToTensor(),<br>
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])</p>
<p>train_dataset = ImageFolder(root=os.path.join(dataset_path , ‘train’), transform=train_transform)<br>
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, num_workers=2)</p>
<p>model = models.inception_v3(pretrained=True, aux_logits=False).to(device)<br>
model.fc = nn.Linear(2048, 2)<br>
print('model on cuda: ', next(model.parameters()).is_cuda)</p>
<p>criterion = nn.CrossEntropyLoss()<br>
optimizer = SGD(model.parameters(), lr=0.005, momentum=0.9)<br>
scaler = GradScaler()</p>
<p>for epoch in range(3):</p>
<pre><code>print('epoch: ', epoch+1)
for i, data in enumerate(train_loader, 0):
    inputs, labels = data
    if torch.cuda.is_available():
        inputs, labels = inputs.to(device),labels.to(device)
    print('inputs on cuda: ', inputs.is_cuda, ', type: ', inputs.dtype)
    print('labels on cuda: ', labels.is_cuda, ', type: ', labels.dtype)
    with autocast():
        output = model(inputs)
        print('output on cuda: ', output.is_cuda, ', type: ', output.dtype)
        loss = criterion(output, labels)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()</code></pre>

username: ptrblck
cooked: <p>Your new module should still be on the CPU since you are:</p>
<ul>
<li>creating the inception_v3` model</li>
<li>push it to the GPU</li>
<li>replace <code>model.fc</code> with a new <code>nn.Linear</code> on the CPU</li>
<li>don’t push it to the device</li>
</ul>

username: Nimrod_Daniel
cooked: <p>Thanks! Didn’t know the replacement occurs only on the CPU. If I make some changes I later move it to the device right after (didn’t do that here), but good to know that any replacement occurs only on the CPU</p>

username: ptrblck
cooked: <p>Yes, the reason for this is because:</p>
<pre><code class="lang-python">model.fc = nn.Linear(2048, 2)
</code></pre>
<p>initializes it with a plain <code>nn.Linear</code> module, which will use the CPU as its default device and you would need to explicitly push the modules to the GPU if needed.<br>
The common approach would be to manipulate your model first, then push the entire model to the GPU once.</p>