edit a few line for test purposes

Author: nurgoni

src/onepiece_classify/trainer/trainer.py

@@ -28,7 +28,7 @@ def __init__(
         self._val_loss = []
         self._val_acc = []
 
-    def train_step(self, batch, batch_idx):
+    def train_step(self, batch, batch_idx: int = None):
         X, y = batch
         output = self.model(X)
 
@@ -44,7 +44,7 @@ def train_step(self, batch, batch_idx):
 
         return train_loss, train_acc
 
-    def val_step(self, batch, batch_idx):
+    def val_step(self, batch, batch_idx: int=None):
         X, y = batch
         output = self.model(X)
 
@@ -56,29 +56,33 @@ def val_step(self, batch, batch_idx):
 
         return val_loss, val_acc
 
-    def train_batch(self, epoch):
+    def train_batch(self, epoch: int=None):
         self.model.train()
 
         self.train_loss = 0
         self.train_acc = 0
         for batch_idx, batch in enumerate(self.loader.train_loader):
-            self.train_loss, self.train_acc += self.train_step(batch=batch, batch_idx=batch_idx)
-        
+            _loss, _acc = self.train_step(batch=batch, batch_idx=batch_idx)
+            self.train_loss += _loss
+            self.train_acc += _acc
+
         avg_loss = self.train_loss/len(self.loader.train_loader)
         avg_acc = self.train_acc/len(self.loader.train_loader)
 
         return avg_loss, avg_acc
 
 
-    def val_batch(self, epoch):
+    def val_batch(self, epoch: int=None):
         self.model.eval()
 
         self.val_loss = 0
         self.val_acc = 0
         with torch.no_grad():
             for batch_idx, batch in enumerate(self.loader.valid_loader):
-                self.val_loss, self.val_acc += self.val_step(batch=batch, batch_idx=batch_idx)
-        
+                _loss, _acc = self.val_step(batch=batch, batch_idx=batch_idx)
+                self.val_loss += _loss
+                self.val_acc += _acc
+
         avg_loss = self.val_loss/len(self.loader.valid_loader)
         avg_acc = self.val_acc/len(self.loader.valid_loader)
 
@@ -104,10 +108,3 @@ def fit(self, model: nn.Module, loader: DataLoader) -> Dict[str, float]:
         self.loader = loader
         self.run()
         # self.save_model()
-
-trainer = Trainer(
-    max_epochs=5, 
-    loss_fn=nn.CrossEntropyLoss(), 
-    optim=torch.optim.Adam()
-    learning_rate=0.001
-)