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)