Update dataset.py

dataset.py

@@ -12,8 +12,11 @@ class CXRMate2Dataset:
         self.history = history
         self.study_id_to_index = dict(zip(self.dataset['study_id'], range(len(self.dataset)), strict=True))
 
-    def __getitem__(self, idx):
-        batch = self.dataset[idx]
+    def __getitem__(self, key):
+        if not isinstance(key, int):
+            return self.dataset[key]
+
+        batch = self.dataset[key]
 
         if 'views' not in batch:
             batch['views'] = [None] * len(batch['images'])
@@ -26,7 +29,7 @@ class CXRMate2Dataset:
 
         if self.history:
             
-            if not batch['prior_study_ids']:
+            if batch['prior_study_ids']:
             
                 # Sort by datetime to ensure correct order:
                 assert all(i is not None and not (isinstance(i, float) and np.isnan(i)) for i in batch['prior_study_datetimes'])