Update tasks/image.py

tasks/image.py

@@ -50,13 +50,13 @@ class SmokeDataset(Dataset):
         image = example["image"]
         annotation = example.get("annotations", "").strip()
 
-        # Resize and preprocess the image directly here
+        # Resize and preprocess the image
         image = image.resize((512, 512))
         image = np.array(image)[:, :, ::-1]  # Convert RGB to BGR
         image = np.array(image, dtype=np.float32) / 255.0
 
-        # Return both the preprocessed image and annotation
-        return torch.tensor(image).permute(2, 0, 1), annotation
+        # Return both the preprocessed image tensor and annotation
+        return torch.tensor(image, dtype=torch.float32).permute(2, 0, 1), annotation
 
 
 def preprocess(image):