Fix how dataset is accessed (#758)

lightly-ai · Apr 5, 2022 · c9fa2c5 · c9fa2c5
1 parent 947d3fd
commit c9fa2c5
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/lightly/embedding/embedding.py b/lightly/embedding/embedding.py
@@ -106,11 +106,12 @@ def embed(self,
         self.model.eval()
         embeddings, labels, filenames = None, None, []
 
+        dataset = dataloader.dataset
         if lightly._is_prefetch_generator_available():
             dataloader = BackgroundGenerator(dataloader, max_prefetch=3)
 
         pbar = tqdm(
-            total=len(dataloader.dataset),
+            total=len(dataset),
             unit='imgs'
         )
 
@@ -157,7 +158,7 @@ def embed(self,
             embeddings = embeddings.cpu().numpy()
             labels = labels.cpu().numpy()
 
-        sorted_filenames = dataloader.dataset.get_filenames()
+        sorted_filenames = dataset.get_filenames()
         sorted_embeddings = sort_items_by_keys(
             filenames, embeddings, sorted_filenames
         )