Data not going to proper device

vision
#1 
model = models.vgg16(pretrained=True).to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.NLLLoss().to(device) # using NLL since our last output is log
BATCH_SIZE = 50
EPOCHS = 1

def load_data():
    mean=[0.485, 0.456, 0.406]
    std=[0.229, 0.224, 0.225]

    data_transforms = {
        'train': transforms.Compose([
            transforms.RandomResizedCrop(224),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize(mean, std)
        ]),
        'val': transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean, std)
        ])
    }

    path = os.path.dirname(__file__)
    data = {
        'train': datasets.ImageFolder(root, data_transforms['train']),
        'val': datasets.ImageFolder(root', data_transforms['val'])
    }
    dataloaders = {
        'train': DataLoader(data['train'], batch_size=BATCH_SIZE, shuffle=True),
        'val': DataLoader(data['val'], batch_size=BATCH_SIZE, shuffle=True)
    }
    return dataloaders

    dataloaders = load_data()
    inputs, classes = next(iter(dataloaders['train']))
    train(inputs, classes, dataloaders['train'])

def train(inputs, classes, dataloader):
    model.train()

    # freeze models weights and unfreeze few layers we want
    for layer in model.parameters():
        layer.requires_grad = False


    # custom classifier on 6th sequence
    model.classifier[6] = nn.Sequential(
        nn.Linear(4096, 256), # in features from prevous out of model.classifier
        nn.ReLU(),
        nn.Dropout(0.4),
        nn.Linear(256, 2),
        nn.LogSoftmax(dim=1)
    )
    # print(model.classifier)

    for epoch in range(EPOCHS):
        for inputs, classes in tqdm(dataloader):
            inputs = inputs.to(device)
            classes = classes.to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, classes)
            loss.backward()
            optimizer.step()

Traceback (most recent call last):
  File "ml\train&test.py", line 114, in <module>
    train(inputs, classes, dataloaders['train'])
  File "ml\train&test.py", line 88, in train
    outputs = model(inputs)
  File "C:\Users\Vijay\Anaconda3\envs\desert\lib\site-packages\torch\nn\modules\module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "C:\Users\Vijay\Anaconda3\envs\desert\lib\site-packages\torchvision\models\vgg.py", line 46, in forward
    x = self.classifier(x)
  File "C:\Users\Vijay\Anaconda3\envs\desert\lib\site-packages\torch\nn\modules\module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "C:\Users\Vijay\Anaconda3\envs\desert\lib\site-packages\torch\nn\modules\container.py", line 100, in forward
    input = module(input)
  File "C:\Users\Vijay\Anaconda3\envs\desert\lib\site-packages\torch\nn\modules\module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "C:\Users\Vijay\Anaconda3\envs\desert\lib\site-packages\torch\nn\modules\container.py", line 100, in forward
    input = module(input)
  File "C:\Users\Vijay\Anaconda3\envs\desert\lib\site-packages\torch\nn\modules\module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "C:\Users\Vijay\Anaconda3\envs\desert\lib\site-packages\torch\nn\modules\linear.py", line 87, in forward
    return F.linear(input, self.weight, self.bias)
  File "C:\Users\Vijay\Anaconda3\envs\desert\lib\site-packages\torch\nn\functional.py", line 1370, in linear
    ret = torch.addmm(bias, input, weight.t())
RuntimeError: Expected object of device type cuda but got device type cpu for argument #1 'self' in call to _th_addmm

I don’t know why in the train function with the loop over inputs and classes are not going to gpu.

#2

I think the error is raised, since model.classifier[0] might not have been pushed to the device.
While you are using model = models.vgg().to(device) correctly, you are manipulating the classifier later without calling .to() again on this part of the model, so it should stay on the CPU.

#3 
    model.classifier[6] = nn.Sequential(
        nn.Linear(4096, 256), # in features from prevous out of model.classifier
        nn.ReLU(),
        nn.Dropout(0.4),
        nn.Linear(256, 2),
        nn.LogSoftmax(dim=1)
    ).to(device)

Missed that .to(device) on the end. Thank you very much, ptrblck.