Update tasks/image.py

tasks/image.py

@@ -1,4 +1,5 @@
 import torch
+import os
 import torch.nn as nn
 import torch.optim as optim
 from torchvision import transforms
@@ -93,10 +94,9 @@ async def evaluate_image(request: ImageEvaluationRequest):
     
     # Load and prepare the dataset
     dataset = load_dataset(request.dataset_name, token=os.getenv("HF_TOKEN"))
-    
+
     # Split dataset
-    train_test = dataset["train"]
-    test_dataset = dataset["val"]
+    test_dataset = dataset["test"]
     
     # Start tracking emissions
     tracker.start()