enable tensor index.

vincentqb · vincentqb · commit d13a4c6c6f1e · 2019-08-16T10:07:21.000-04:00
diff --git a/beginner_source/data_loading_tutorial.py b/beginner_source/data_loading_tutorial.py
@@ -136,6 +136,9 @@ def __len__(self):
         return len(self.landmarks_frame)
 
     def __getitem__(self, idx):
+        if torch.is_tensor(idx):
+            idx = idx.tolist()
+
         img_name = os.path.join(self.root_dir,
                                 self.landmarks_frame.iloc[idx, 0])
         image = io.imread(img_name)
