Optimize DataLoader iteration in WrappedDataLoader (#2375)

pytorch · Jun 1, 2023 · e1ec4bd · e1ec4bd
1 parent e2a7ab0
commit e1ec4bd
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/beginner_source/nn_tutorial.py b/beginner_source/nn_tutorial.py
@@ -795,8 +795,7 @@ def __len__(self):
         return len(self.dl)
 
     def __iter__(self):
-        batches = iter(self.dl)
-        for b in batches:
+        for b in self.dl:
             yield (self.func(*b))
 
 train_dl, valid_dl = get_data(train_ds, valid_ds, bs)