Bugfix: meta_train_dataset etc. are attributes of benchmark

tristandeleu · Apr 10, 2020 · f0783b5 · f0783b5
1 parent 3c0ab81
commit f0783b5
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__':