Adafactor: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! about pytorch_optimizer HOT 3 CLOSED

from timm import create_model
from tqdm import tqdm

import torch
from torch.nn import functional as F

from pytorch_optimizer import load_optimizer

bs = 16
backbone = 'resmlp_12_distilled_224'

model = create_model(backbone, pretrained=False, num_classes=1)
model.train()
model.cuda()

optimizer = load_optimizer('adafactor')(model.parameters(), warmup_init=True, weight_decay=1e-3)

inp = torch.randn((bs, 3, 224, 224), dtype=torch.float32)
y = torch.randn((bs, 1), dtype=torch.float32)

inp = inp.cuda()
y = y.cuda()

for _ in tqdm(range(50)):
    optimizer.zero_grad()

    F.binary_cross_entropy_with_logits(model(inp), y).backward()

    optimizer.step()