Не работает cuda в python на видеокарте NVIDIA (Windows 10)

Обучаю модель для сегментации картинок, обычно запускал на ЦП и время одной эпохи было 40 минут. Решил использовать CUDA. На видеокарте стоит 12.7, установил для Python последнюю версию что есть (12.4). И скорость вычислений не увеличилась, хотя устройство показано в программе что доступно, но толку ноль. Помогите. Использую библиотеку PyTorch для взаимодействия с видеокартой. Примерный код

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

num_classes = len(reverse_class_dict)
model = UNet(in_channels=3, out_channels=num_classes).to(device)
criterion = nn.CrossEntropyLoss()
if (torch.cuda.device_count()>1):
    print("Используется", torch.cuda.device_count(), "ГПУ")
    model = DataParallel(model)
else:
    print("Доступен только один графический процессор")
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

def train_model(model, dataloader, criterion, optimizer, num_epochs=10):
    model.train()
    scaler = torch.GradScaler(device="cuda")
    for epoch in range(num_epochs):
        epoch_loss = 0
        total_pixels, correct_pixels = 0, 0

        for batch in tqdm(dataloader, desc=f"Epoch {epoch + 1}/{num_epochs}"):
            images = batch['image'].to(device)
            masks = batch['mask'].to(device)

            optimizer.zero_grad()
            with torch.autocast(device_type="cuda"):
                outputs = model(images)
                loss = criterion(outputs, masks)

            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
            epoch_loss += loss.item()
            preds = torch.argmax(outputs, dim=1)
            total_pixels += masks.numel()
            correct_pixels += (preds == masks).sum().item()

        accuracy = correct_pixels / total_pixels
        print(f"Epoch {epoch + 1}/{num_epochs}, Loss: {epoch_loss:.4f}, Accuracy: {accuracy:.2f}")