diff options
Diffstat (limited to 'train_ti.py')
| -rw-r--r-- | train_ti.py | 4 |
1 files changed, 3 insertions, 1 deletions
diff --git a/train_ti.py b/train_ti.py index 83043ad..dd015f9 100644 --- a/train_ti.py +++ b/train_ti.py | |||
| @@ -774,7 +774,9 @@ def main(): | |||
| 774 | num_train_epochs = args.num_train_epochs | 774 | num_train_epochs = args.num_train_epochs |
| 775 | sample_frequency = args.sample_frequency | 775 | sample_frequency = args.sample_frequency |
| 776 | if num_train_epochs is None: | 776 | if num_train_epochs is None: |
| 777 | num_train_epochs = math.ceil(args.num_train_steps / len(datamodule.train_dataset)) | 777 | num_train_epochs = math.ceil( |
| 778 | args.num_train_steps / len(datamodule.train_dataset) | ||
| 779 | ) * args.gradient_accumulation_steps | ||
| 778 | sample_frequency = math.ceil(num_train_epochs * (sample_frequency / args.num_train_steps)) | 780 | sample_frequency = math.ceil(num_train_epochs * (sample_frequency / args.num_train_steps)) |
| 779 | 781 | ||
| 780 | optimizer = create_optimizer( | 782 | optimizer = create_optimizer( |