code/train_alexnet.py
@@ -179,7 +179,7 @@ loss_roll += alpha * loss_reg_roll loss_seq = [loss_yaw, loss_pitch, loss_roll] grad_seq = [torch.Tensor(1).cuda(gpu) for _ in range(len(loss_seq))] grad_seq = [torch.ones(1).cuda(gpu) for _ in range(len(loss_seq))] torch.autograd.backward(loss_seq, grad_seq) optimizer.step()
