Automatic Mixed Precision: Input type (torch.cuda.HalfTensor) and weight type (torch.FloatTensor) should be the same

username: Nimrod_Daniel
cooked: <p>Hi,<br>
I’m trying to train using amp, but unfortunately I get the:<br>
Input type (torch.cuda.HalfTensor) and weight type (torch.FloatTensor) should be the same<br>
error. I’ve seen others had this problem, but only because they didn’t use autocast correctly, which is not the case here. I created a following shorter code that reproduce the error:<br>
(torch 1.10.2, python 3.6.9, T4 GPU)</p>
<p>import torch</p>
<p>import torch.nn as nn<br>
import torch.nn.functional as F<br>
from torch.optim import SGD<br>
from torchvision import transforms,models<br>
import torchvision.datasets as datasets</p>
<p>from torch.cuda.amp import GradScaler, autocast</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(32),<br>
transforms.RandomHorizontalFlip(p=.50),<br>
transforms.ToTensor(),<br>
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])</p>
<p>test_transform = transforms.Compose([<br>
transforms.Resize(32),<br>
transforms.ToTensor(),<br>
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])</p>
<p>traindata = datasets.CIFAR10(root=‘.’, train=True,download=True, transform=train_transform)<br>
train_loader = torch.utils.data.DataLoader(traindata, batch_size=96,shuffle=True)</p>
<p>class Net(nn.Module):</p>
<pre><code>def __init__(self):
    super().__init__()
    self.conv1 = nn.Conv2d(3, 6, 5)
    self.pool = nn.MaxPool2d(2, 2)
    self.conv2 = nn.Conv2d(6, 16, 5)
    self.fc1 = nn.Linear(16 * 5 * 5, 120)
    self.fc2 = nn.Linear(120, 84)
    self.fc3 = nn.Linear(84, 10)

def forward(self, x):
    x = self.pool(F.relu(self.conv1(x)))
    x = self.pool(F.relu(self.conv2(x)))
    x = torch.flatten(x, 1) # flatten all dimensions except batch
    x = F.relu(self.fc1(x))
    x = F.relu(self.fc2(x))
    x = self.fc3(x)
    return x
</code></pre>
<p>model = Net()</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(5):</p>
<pre><code>for i, data in enumerate(train_loader, 0):

    inputs, labels = data
    inputs, labels = inputs.to(device),labels.to(device)
            
    with torch.cuda.amp.autocast():
        output = model(inputs)
        loss = criterion(output, labels)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()
</code></pre>
<p>Plus,in my code (just different dataloader, that build using ImageFolder), if I use inception rather than a similiar net from above (2 outputs instead):</p>
<pre><code>    model = models.inception_v3(pretrained=True, aux_logits=False).to(device)
    model.fc = nn.Linear(2048, 2)
</code></pre>
<p>Then I get:</p>
<blockquote>
<p>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)</p>
</blockquote>

username: ptrblck
cooked: <p>You’ve forgot to push the model to the device via <code>model.to(device)</code>.<br>
Adding it fixes the device mismatch.</p>

username: Nimrod_Daniel
cooked: <p>Right, when I created the simple code I forgot to push it to the device <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<p>Anyway, any idea why would I get an error such as:</p>
<blockquote>
<p>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)</p>
</blockquote>
<p>When I have a similar training block (with inception for instance) and the model is pushed to the GPU?</p>

username: Nimrod_Daniel
cooked: <p>I think it’d be better if I provide the code (in a new question)</p>