unteste update using troch-trainer

1 files changed, 2 insertions, 3 deletions

diff --git a/autoencoder.py b/autoencoder.py
index 58326fe..23c47f9 100644
--- a/autoencoder.py
+++ b/autoencoder.py
@@ -5,7 +5,7 @@ from torchvision import datasets, transforms
 
 import matplotlib.pyplot as plt
 
-from trainer import Trainer
+from trainer import train
 
 
 class PrintLayer(torch.nn.Module):
@@ -109,12 +109,11 @@ def main():
     model.to(device)
 
     # Train model
-    trainer = Trainer()
     lr = 1e-3
     epochs = 1
     optimizer = torch.optim.Adam(model.parameters(), lr=lr)
     criterion = torch.nn.functional.binary_cross_entropy
-    trainer.train(model, dataloader, epochs, optimizer, criterion)
+    train(model, dataloader, epochs, optimizer, criterion)
 
     # Save model
     torch.save(model.state_dict(), 'autoencoder.pth')