fix trainer

main.py

@@ -123,6 +123,7 @@ if __name__ == "__main__":
         # save_steps=5,
         # eval_steps=5,
     )
+    trainer = None
 
     # PretrainedConfig.from_json_file(os.path.join(resume_from_checkpoint, CONFIG_NAME))
     last_checkpoint_path = None
@@ -163,16 +164,20 @@ if __name__ == "__main__":
                                              )
         test_dataset = test_dataset.shard(num_test_sub_shard, idx_sub_shard)
         # Init trainer
-        trainer = Trainer(
-            model=w2v_ctc_model,
-            data_collator=data_collator,
-            args=training_args,
-            compute_metrics=compute_metrics_fn(w2v_ctc_processor),
-            train_dataset=train_dataset,
-            eval_dataset=test_dataset,
-            tokenizer=w2v_ctc_processor.feature_extractor,
-            callbacks=[BreakEachEpoch()]  # Manual break end of epoch because each epoch loop over a shard
-        )
+        if trainer is None:
+            trainer = Trainer(
+                model=w2v_ctc_model,
+                data_collator=data_collator,
+                args=training_args,
+                compute_metrics=compute_metrics_fn(w2v_ctc_processor),
+                train_dataset=train_dataset,
+                eval_dataset=test_dataset,
+                tokenizer=w2v_ctc_processor.feature_extractor,
+                callbacks=[BreakEachEpoch()]  # Manual break end of epoch because each epoch loop over a shard
+            )
+        else:
+            trainer.train_dataset = train_dataset
+            trainer.eval_dataset = test_dataset
 
         logging.get_logger().info('Train shard idx: {} / {}'.format(train_dataset_shard_idx + 1, num_train_shards))
         logging.get_logger().info(