Update chess-mamba-vs-xformer/train_bygame.py

chess-mamba-vs-xformer/train_bygame.py

@@ -148,7 +148,9 @@ ctx = nullcontext() if device_type == 'cpu' else torch.amp.autocast(device_type=
 # poor man's data loader
 data_dir = os.path.join('data', dataset)
 current_train_file_index = 0
-train_files = glob.glob(os.path.join(data_dir, 'train*.parquet'))
+train_files = glob.glob(os.path.join(data_dir, 'train*.parquet')) + \
+              glob.glob(os.path.join(data_dir, 'stable*.parquet')) + \
+              glob.glob(os.path.join(data_dir, 'anneal*.parquet'))
 train_datasets = []
 for f in train_files:
     dataset = pq.read_table(f).to_pandas()