Вычисления происходят не на gpu, хотя torch.cuda.is_available() == True

#device становится device(type='cuda')
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
#тензоры отправляются на gpu (например, y_test_tensor.device == device(type='cuda', index=0))
sc = StandardScaler()
X_train_std = (torch.from_numpy(sc.fit_transform(X_train).astype(np.float32))).to(device)
X_test_std = (torch.from_numpy(sc.transform(X_test).astype(np.float32))).to(device)
y_train_tensor = (torch.tensor(y_train.values, dtype=torch.float32)).to(device)
y_test_tensor = (torch.tensor(y_test.values, dtype=torch.float32)).to(device)
 
dataset = TensorDataset(X_train_std, y_train_tensor)
dataLoader = DataLoader(dataset, batch_size, shuffle=True)
 
#модель тоже перевожу на gpu (next(model_base.parameters()).device == device(type='cuda', index=0))
model_base = nn.Sequential(
    nn.Linear(inp_size, 20),
    nn.ReLU(),
    nn.Dropout(p=0.5),
    nn.Linear(20, 1),
    nn.Sigmoid()
).to(device)
loss_fn = nn.BCELoss()
optimizer = torch.optim.Adam(model_base.parameters(), weight_decay=1e-4)
 
def train(model, num_epochs, train_dl, x_valid, y_valid):
    for epoch in range(num_epochs):
        model.train()
        for x_batch, y_batch in train_dl:
            pred = model(x_batch)[:, 0]
            loss = loss_fn(pred, y_batch)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
 
        model.eval()
        
train(model_base, num_epochs, dataLoader, X_test_std, y_test_tensor)

print("x_batch: ")
print(x_batch.device)
print("pred: ")
print(pred.device)