FlowModel / core /trainer.py
LunaStev
GitHub to Huggingface
cff2458
raw
history blame contribute delete
673 Bytes
import torch.optim as optim
import torch.nn as nn
class Trainer:
def __init__(self, model):
self.model = model
self.criterion = nn.CrossEntropyLoss()
self.optimizer = optim.SGD(self.model.parameters(), lr=0.01)
def train(self, train_loader, epochs=5):
for epoch in range(epochs):
for images, labels in train_loader:
self.optimizer.zero_grad()
outputs = self.model(images.view(-1, 28*28))
loss = self.criterion(outputs, labels)
loss.backward()
self.optimizer.step()
print(f"Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}")