@ptrblck here is my full code:
import torch
import torchvision
import torchvision.transforms as transforms

transform=transforms.Compose([transforms.Grayscale(3),transforms.Resize(256),transforms.ToTensor(),transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))])

trainset=torchvision.datasets.MNIST(root='./data',train=True,download=True,transform=transform)
trainloader=torch.utils.data.DataLoader(trainset,batch_size=32,shuffle=True,num_workers=2)

testset=torchvision.datasets.MNIST(root='./data',train=False,download=True,transform=transform)
testloader=torch.utils.data.DataLoader(testset,batch_size=32,shuffle=False,num_workers=2)

import torch.nn as nn
import torch.nn.functional as F

device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

import torchvision.models as models
alexnet=models.alexnet(pretrained=True)
alexnet.classifier[6]=nn.Linear(4096,10)
print(alexnet)
print('Model Downloaded')

class MnistAlexnet(nn.Module):

        def __init__(self,alexnet):
            super(MnistAlexnet,self).__init__()
            self.alexnet=alexnet

        def forward(self,x):
            return self.alexnet(x)

model=MnistAlexnet(alexnet)  
print(model)


import time
import torch.optim as optim

batch_size=128
epochs=5
model.to(device)
criterion=nn.CrossEntropyLoss()
optimizer=optim.Adam(alexnet.parameters(),lr=0.001)


for epoch in range(epochs):
        running_loss=0.0
        total=0
        probs=0
        correct_classified=0
        start_time=time.time()
        model.train()
        for i,data in enumerate(trainloader):
            inputs,labels=data
            inputs,labels=inputs.to(device),labels.to(device)
            optimizer.zero_grad()
            outputs=model(inputs)
            loss=criterion(outputs,labels)
            loss.backward()
            optimizer.step()
           # _,predicted=torch.max(outputs.data,1)
            predicted=torch.argmax(outputs,1)
            total+=labels.size(0)
            probs=F.softmax(outputs,1)[1]
            #print("probs=",probs)
            correct_classified+=(predicted==labels).sum().item()
            running_loss+=loss.item()
            if i % 200 == 199:
               avg_loss=running_loss/200
               print('Epoch:[%d, %5d] Batch: %5d loss: %.3f' % (epoch+1,i+1,i+1,avg_loss))
            running_loss=0.0
            train_acc=(100*correct_classified/total)
       # print("Time/epoch: {} sec".format(time.time()-start_time))
       # train_acc=(correct_classified/total)
       # print('Train Accuracy :%d'%(train_acc))
        c=0
        total=0
        model.eval()
        with torch.no_grad():
            for data in testloader:
               images,labels=data
               inputs,labels=images.to(device),labels.to(device)
               outputs=model(inputs)
              # _,predicted=torch.max(outputs.data,1)
               predicted=torch.argmax(outputs,1)
               total+=labels.size(0)
               c+=(predicted==labels).sum().item()
               test_acc=(100*c/total)

        print("c=",c," total=",total)
        print('Accuracy of the network on test images:%d %%' % test_acc)



print('trained')