From 9d6d3f3c7f8f01e24d635c3f5d9b43fa697d5f9f Mon Sep 17 00:00:00 2001
From: gdamms <damguillotin@gmail.com>
Date: Tue, 2 Jul 2024 11:39:08 +0200
Subject: unteste update using troch-trainer

---
 autoencoder.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

(limited to 'autoencoder.py')

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')
-- 
cgit v1.3.1