[HOTFIX] Convert target to Tensor in ClassificationDataset (#208)

v7labs · Sep 22, 2021 · b4379bd · b4379bd
1 parent f46cc8e
commit b4379bd
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/darwin/torch/dataset.py b/darwin/torch/dataset.py
@@ -126,7 +126,7 @@ def get_target(self, index: int):
                 idx = self.classes.index(tag)
                 target[idx] = 1
         else:
-            target = self.classes.index(tags[0])
+            target = torch.tensor([self.classes.index(tags[0])])
 
         return target