Hello Everyone,
I am trying to train a model that learns to predict a single value(float) given an image using the pretrained resnet18 model. I am using the MSE loss but I am getting a Runtime error below.
RuntimeError: Found dtype Double but expected Float. from loss.backward()
Below is my code:
resnet18= models.resnet18(pretrained=True)
num_ftrs = resnet18.fc.in_features
resnet18.fc = nn.Linear(num_ftrs , 1)
def train_model(model, dataloaders, device, num_epochs=2, is_train=True):
acc_history=[]
loss_history=[]
best_acc= 0.0
for epoch in tqdm(range(num_epochs)):
running_loss = 0.0
running_corrects=0
#Iterate over dataloader
for x,y in dataloaders:
for image,target in zip(x,y):
image = image.to(device)
target = target.to(device)
model.to(device)
#zero the parameter gradients
optimizer.zero_grad()
criterion = nn.MSELoss()
# forward
out = model(image.unsqueeze(0))
print(out)
loss = criterion(out, target)
_,preds =torch.max(out,1)
#backward
loss.backward()
optimizer.step()