Function returned an invalid gradient at index 0

1

I copied the pytorch code for concrete dropout from the authors github repo and I put the model on a GPU in my own code by calling model.to(device).

I now get the error RuntimeError: Function DivBackward0 returned an invalid gradient at index 0 - expected type torch.FloatTensor but got torch.cuda.FloatTensor and I get that somewhere in the model there is a tensor that is not on the GPU, but I can’t find it and I’ve been printing out everything in the model trying to find out where it is.

Is there an easy way to figure out where this is being caused?

2

Is the loss Maybe calculated on CPU ??

Try to run your code on cpu to see if you still receive the error

3

I don’t think so, below is the loss function and the calculation

def heteroscedastic_loss(true, mean, log_var):
    precision = torch.exp(-log_var)
    return torch.mean(torch.sum(precision * (true - mean) ** 2 + log_var, 1), 0)

...

loss = heteroscedastic_loss(y, mean, log_var) + regularization
loss.backward()

and I tried it on the cpu and everything works fine