I am working on the MNIST dataset for an assignment and it seems to me I am stuck at some point for long. I have written my code for LogisticRegression and when I try to train the model it is not working as expected but instead, it throws an error that says:
Input tensor should be a torch tensor. Got <class 'PIL.Image.Image'>.
Dataloader:
train_dataloader = torch.utils.data.DataLoader(mnist_train, batch_size, shuffle=True)
My Code:
class LogisticRegression(torch.nn.Module):
def __init__(self, input_dim, output_dim):
super(LogisticRegression, self).__init__()
self.linear = torch.nn.Linear(input_dim, output_dim)
def forward(self, x):
x = x.reshape(-1, 784)
outputs = self.linear(x)
return outputs
Train the model:
for epoch in range(epochs):
for i, (images, labels) in enumerate(train_dataloader):
optimizer.zero_grad()
y_pred = model(images)
loss = criterion(y_pred, labels)
loss.backward()
optimizer.step()
Where am I going wrong with the code?