Fix no data bug with --div-data

This commit is contained in:
Yin Li 2020-05-07 21:08:15 -04:00
parent 8bb6e85c6d
commit 9b456d6b1a

View File

@ -55,6 +55,8 @@ class FieldDataset(Dataset):
self.in_files = self.in_files[rank * files : (rank + 1) * files]
self.tgt_files = self.tgt_files[rank * files : (rank + 1) * files]
assert len(self.in_files) > 0, 'files not divisible among ranks'
self.in_chan = [np.load(f).shape[0] for f in self.in_files[0]]
self.tgt_chan = [np.load(f).shape[0] for f in self.tgt_files[0]]