I don’t know how flatten the image during training for MNIST which is (1 * 28 * 28) i was doing image.view(-1,28 * 28) how about CIFA10 which is (3* 32*32) my code on class and training are
class LogisticRegressionCIFAR(nn.Module):
def init(self,input_dim, output_dim):
super (LogisticRegressionCIFAR, self).init()
self.linear = nn.Linear(input_dim,output_dim)
def forward(self,x):
    out = self.linear(x)
    return out

and for  training is
iter = 0
for epoch in range (5):
for i, (images,labels) in enumerate(train_loader):
images = Variable(images.view(-3,32*32))
labels = Variable(labels)
    optimizer.zero_grad()
    outputs = model(images)
    loss = criterion(outputs,labels)
    loss.backward()
    optimizer.step()
    
    iter +=1
    
    if iter % 500 ==0:
        correct = 0
        total = 0
        for images,labels in test_loader:
            images = Varaible(images.view(-3,32*32))
            
            outputs = model(images)
            
            _,predicted = torch.max(outputs.data,1)
            total += labels.size(0)
            correct += (predicted == labels).sum()
            
        accuracy = 100 * correct / total
        
        print ('Iteration: {}. Loss {} . Accuracy {}'.format(iter, loss.data[0]), accuracy)