flpelerin commited on
Commit
7f3b4e4
·
1 Parent(s): 52a257c

Update 2 files

Browse files

- /trainer.cli.py
- /dataset.py

Files changed (2) hide show
  1. dataset.py +2 -0
  2. trainer.cli.py +1 -1
dataset.py CHANGED
@@ -1,4 +1,6 @@
1
  from datasets import load_dataset
 
 
2
 
3
  from util import Config, GetDevice
4
 
 
1
  from datasets import load_dataset
2
+ import numpy as np
3
+ import torch
4
 
5
  from util import Config, GetDevice
6
 
trainer.cli.py CHANGED
@@ -37,7 +37,7 @@ if __name__ == '__main__':
37
  batches, num_batches = dataset.batch(ids)
38
 
39
 
40
- print(f"dataset ids: {dataset.ids}")
41
 
42
 
43
  trainer = Trainer(config)
 
37
  batches, num_batches = dataset.batch(ids)
38
 
39
 
40
+ print(f"batches: {num_batches}")
41
 
42
 
43
  trainer = Trainer(config)