Tensor for 'out' is on CPU, Tensor for argument #1 'self' is on CPU, but expected them to be on GPU (while checking arguments for addmm)

Hi all,

I tried to solve simple code based on MNIST

import numpy as np
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
from torchvision import datasets, transforms
dev = 'cuda'
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'

batch_size = 1
test_batch_size = 1
train_loader = torch.utils.data.DataLoader (datasets.MNIST('dataset/', train=True, download=True,
                                                           transform= transforms.Compose([transforms.ToTensor(),
                                                            ])),
                                                            batch_size = batch_size, shuffle=True)

image, label = next(iter(train_loader))
image = image.to(dev)

with torch.no_grad(): 

    flatten = image.view(1, 28 * 28).to(dev) 
    print(flatten.shape)

    lin = nn.Linear(784, 10)(flatten) 
    print(lin)

Error occurs at lin = nn.Linear (784, 10)(flatten) even though I designated ‘flatten’ to ‘cuda’.

Traceback (most recent call last):
  File "C:/Users/yoonh/Works/main.py", line 59, in <module>
    lin = nn.Linear(784, 10)(flatten) # 784를 받아서 10개로 내보냄
  File "C:\ProgramData\Anaconda3\lib\site-packages\torch\nn\modules\module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "C:\ProgramData\Anaconda3\lib\site-packages\torch\nn\modules\linear.py", line 94, in forward
    return F.linear(input, self.weight, self.bias)
  File "C:\ProgramData\Anaconda3\lib\site-packages\torch\nn\functional.py", line 1753, in linear
    return torch._C._nn.linear(input, weight, bias)
RuntimeError: Tensor for 'out' is on CPU, Tensor for argument #1 'self' is on CPU, but expected them to be on GPU (while checking arguments for addmm)

I am looking forward to see any help to solve this error.

Kind regards,
Yoonho

Double post from here.