Merge pull request #6 from ferhah/bugfix_dataset_close

Bugfix: meta_train_dataset etc. are attributes of benchmark
tristandeleu · Apr 11, 2020 · c7d994a · c7d994a
2 parents 3c0ab81 + f0783b5
commit c7d994a
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/train.py b/train.py
@@ -90,9 +90,9 @@ def main(args):
             with open(args.model_path, 'wb') as f:
                 torch.save(benchmark.model.state_dict(), f)
 
-    if hasattr(meta_train_dataset, 'close'):
-        meta_train_dataset.close()
-        meta_val_dataset.close()
+    if hasattr(benchmark.meta_train_dataset, 'close'):
+        benchmark.meta_train_dataset.close()
+        benchmark.meta_val_dataset.close()
 
 
 if __name__ == '__main__':