Thank you, I’ll continue the discussion here. This works both forwards/backwards on CPU:
import torch

model = torch.nn.Linear(10,10)
x = torch.ones(10,10).detach()
x[0,0] = x[0,0]+float('NaN')

optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
model.zero_grad()
y = model(x)
y[y!=y] = 0 
loss = y.sum()
loss.backward()
optimizer.step()

Could someone verify if it works on GPU?
Also, anyone has insights on the computational cost vs the ol’ ReLU-hack?