add TODO

gsbDBI · Jun 26, 2023 · a4e228c · a4e228c
1 parent aae9751
commit a4e228c
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/bemb/model/bemb_flex_lightning.py b/bemb/model/bemb_flex_lightning.py
@@ -172,6 +172,7 @@ def section_print(input_text):
         test = create_data_loader(dataset_list[2], batch_size=batch_size // 10, shuffle=False, num_workers=num_workers)
 
         section_print('train the model')
+        # TODO: need to change this.
         trainer = pl.Trainer(gpus=1 if ('cuda' in str(self)) else 0,  # use GPU if the model is currently on the GPU.
                             max_epochs=num_epochs,
                             check_val_every_n_epoch=1,