From f576a0d789cf49a15eb2d15647fd7381729a3b32 Mon Sep 17 00:00:00 2001 From: Falko Habel Date: Sat, 22 Feb 2025 10:43:10 +0100 Subject: [PATCH] fixed autocast --- src/aiunn/finetune.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/aiunn/finetune.py b/src/aiunn/finetune.py index 1f18e4d..13cef3e 100644 --- a/src/aiunn/finetune.py +++ b/src/aiunn/finetune.py @@ -15,7 +15,7 @@ class UpscaleDataset(Dataset): combined_df = pd.DataFrame() for parquet_file in parquet_files: # Load data with chunking for memory efficiency - df = pd.read_parquet(parquet_file, columns=['image_512', 'image_1024']).head(5000) + df = pd.read_parquet(parquet_file, columns=['image_512', 'image_1024']).head(2500) combined_df = pd.concat([combined_df, df], ignore_index=True) # Validate data format @@ -126,7 +126,7 @@ for epoch in range(num_epochs): optimizer.zero_grad() # Use automatic mixed precision context - with autocast(device_type=device): + with autocast(device_type="cuda"): outputs = model(low_res) loss = criterion(outputs, high_res)