Do not override optimizer_step (#44)

aitextgen/train.py

@@ -86,20 +86,8 @@ def configure_optimizers(self):
             num_training_steps=self.hparams["num_steps"],
         )
 
-        self.opt = optimizer
-        self.lr_scheduler = scheduler
         return [optimizer], [scheduler]
 
-    def optimizer_step(
-        self, epoch, batch_idx, optimizer, optimizer_idx, second_order_closure=None
-    ):
-        if self.hparams["tpu"]:
-            xm.optimizer_step(optimizer, barrier=True)
-        else:
-            optimizer.step()
-        optimizer.zero_grad()
-        self.lr_scheduler.step()
-
 
 class ATGProgressBar(ProgressBarBase):
     """A variant progress bar that works off of steps and prints periodically."""

setup.py

@@ -3,7 +3,7 @@
 setup(
     name="aitextgen",
     packages=["aitextgen"],  # this must be the same as the name above
-    version="0.2.2",
+    version="0.2.3",
     description="A robust Python tool for text-based AI training and generation using GPT-2.",
     long_description=open("README.md", "r", encoding="utf-8").read(),
     long_description_content_type="text/markdown",