From 81ceddff3b4470284d62981fb98a1055f3a3a6d1 Mon Sep 17 00:00:00 2001 From: Falko Habel Date: Fri, 21 Feb 2025 22:01:54 +0100 Subject: [PATCH] addded device gpu and fixed autocast --- src/aiunn/finetune.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/aiunn/finetune.py b/src/aiunn/finetune.py index 6ea2bdb..1f18e4d 100644 --- a/src/aiunn/finetune.py +++ b/src/aiunn/finetune.py @@ -8,7 +8,7 @@ import csv from tqdm import tqdm import base64 from torch.amp import autocast, GradScaler - +import torch class UpscaleDataset(Dataset): def __init__(self, parquet_files: list, transform=None): @@ -87,7 +87,7 @@ pretrained_model_path = "/root/vision/AIIA/AIIA-base-512" # Load the model using the AIIA.load class method (the implementation copied in your query) model = AIIABase.load(pretrained_model_path) -device = 'cpu' #torch.device("cuda" if torch.cuda.is_available() else "cpu") +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = model.to(device) from torch import nn, optim from torch.utils.data import DataLoader @@ -126,7 +126,7 @@ for epoch in range(num_epochs): optimizer.zero_grad() # Use automatic mixed precision context - with autocast(): + with autocast(device_type=device): outputs = model(low_res) loss = criterion(outputs, high_res)