I modified your code to this in the for loop:
        label=torch.tensor([[0, 0, 0], [0, 1, 0], [0, 0, 2]])
		output=torch.randn(3, 3, requires_grad=True)
		loss=criterion(output, label.float())

And it works.
If you want to do classification you need to use cross entropy loss. Remember to check shape of tensors you pass into the loss. The complete code for mnist is this:
import time
import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torchvision.datasets as dsets
from torch.autograd import Variable
from torchvision.utils import save_image

GPU = 0

def to_img(x):
	x = x.view(x.size(0), 1, 28, 28)
	return x

train_dataset = dsets.MNIST(root='./data',
	train=True,
	transform=transforms.ToTensor(),
	download=True)

test_dataset = dsets.MNIST(root='./data',
	train=False,
	transform=transforms.ToTensor())

batch_size =64
n_iters = 71800

num_epochs = n_iters / (len(train_dataset) / batch_size)
num_epochs = int(num_epochs)

train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
		batch_size=batch_size,
		shuffle=True)

test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
		batch_size=batch_size,
		shuffle=False)

class CNNModel(nn.Module):
	def  __init__(self):
		super().__init__()
		self.cnn1 = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=5, stride=1, padding=2)
		self.relu1 = nn.ReLU()
		self.maxpool1 = nn.MaxPool2d(kernel_size=2)
		
		self.cnn2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=5, stride=1, padding=2)
		self.relu2 = nn.ReLU()
		
		self.maxpool2 = nn.MaxPool2d(kernel_size=2)
		
		self.fc1 = nn.Linear(32 * 7 * 7, 10)

	def forward(self, x):
	    out = self.cnn1(x)
	    out = self.relu1(out)
	   
	    out = self.maxpool1(out)
	   
	    out = self.cnn2(out)
	    out = self.relu2(out)

	    out = self.maxpool2(out)

	    out = out.view(out.size(0), -1)

	    out = self.fc1(out)
	   
	    return out


model = CNNModel()

if torch.cuda.is_available():
	torch.cuda.set_device(GPU)
	model.cuda()
	print("cuda OK")

criterion = nn.CrossEntropyLoss()

best= 12311242

learning_rate = 1e-3

optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=0)

iter = 0
start = time.time()
for epoch in range(2):
	for data in train_loader:
		img, label= data
		label = label.cuda()
		img = img.cuda()
		# nova=to_img(img)

    	# st = time.time()
		optimizer.zero_grad()
		output = model(img)
		
		# label=torch.tensor([[0, 0, 0], [0, 1, 0], [0, 0, 2]])
		# output=torch.randn(3, 3, requires_grad=True)

		loss=criterion(output, label.long())
		loss.backward()
    	
		optimizer.step()
		iter += 1

print('done')