Both tensors are on same device but raises an error both aren't

[ThiruRJST]

"""This module does training and validation of the model."""

import random

import numpy as np
import torch
import torchmetrics
import wandb
from torch import nn
from torch.cuda.amp import GradScaler, autocast
from torch.utils.data import DataLoader
from tqdm import tqdm

from constants import BATCHSIZE, EPOCHS
from preprocess import AudioDataset
from utils.build_model import CNN14_mod
from utils.callbacks import EarlyStopping


def seed_everything(seed=2021):
    """A function that seeds random states to produce reproducible results."""
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.benchmark = True
    torch.backends.cudnn.deterministic = True


def train_loop(epoch, dataloader, scaler, model, loss_fn, optim, device='cuda:0'):
    """Train loop that trains the model."""
    model.train()
    epoch_loss = 0
    metric = torchmetrics.Accuracy()
    step = None

    pbar = tqdm(enumerate(dataloader), total=len(dataloader))
    for step, (img, label) in pbar:
        optim.zero_grad()

        img = img.to(device).float()
        label = label.to(device).long()

        with autocast():
            yhat = model(img)
            # Loss Calculation
            train_loss = loss_fn(input=yhat, target=label)
            out = torch.softmax(yhat, dim=1)

        train_step_acc = metric(out, label)

        scaler.scale(train_loss).backward()
        scaler.step(optim)
        scaler.update()

        epoch_loss += train_loss.item()

    train_epoch_loss = epoch_loss / len(dataloader)
    train_acc = metric.compute()

    wandb.log({'Train_Epoch_Loss': train_epoch_loss, 'Epoch': epoch})
    wandb.log({'Train_Epoch_Acc': train_acc, 'Epoch': epoch})
    wandb.log({'Train_Step_Acc': train_step_acc, 'Step': step})

    metric.reset()

    return train_epoch_loss, train_acc


def val_loop(epoch, dataloader, model, loss_fn, device='cuda:0'):
    """Validation loop to validate the trained model and save it based on performance."""
    model.eval()
    val_epoch_loss = 0
    metric = torchmetrics.Accuracy()
    pbar = tqdm(enumerate(dataloader), total=len(dataloader))
    step = None

    with torch.no_grad():
        for step, (img, label) in pbar:
            img = img.to(device).float()
            label = label.to(device).long()

            yhat = model(img)
            val_loss = loss_fn(input=yhat, target=label)
            out = torch.softmax(yhat, dim=1)

            val_epoch_loss += val_loss.item()

            val_step_acc = metric(out, label)

        val_lossd = val_epoch_loss / len(dataloader)
        val_acc = metric.compute()

        wandb.log({'Val_Epoch_Loss': val_epoch_loss, 'Epoch': epoch})
        wandb.log({'Val_Epoch_Acc': val_acc, 'Epoch': epoch})
        wandb.log({'Val_Epoch_Acc': val_step_acc, 'Step': step})

        metric.reset()

        return val_lossd, val_acc

I'm using torchmetrics to calculate the accuracy of model in pytorch but it raises an error that the tensors aren't in the same device. Actually both the tensors are in same device.

ERROR: RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

Is this a bug or have I made any mistakes?

[SkafteNicki]

Copy of #608 (comment)

The solution is to move your metrics to the correct device (similar to you moving model and data). So instead of

metric = torchmetrics.Accuracy()

do

metric = torchmetrics.Accuracy().to(device)

[ThiruRJST]

Thanks ,it worked