Solved by LeviViana in post #2 
              
              
                I guess it might come from the fact that t and f don’t require grad. 
import torch
prob = torch.randn(5)
prob.requires_grad_(True)
t = torch.ones_like(prob)
f = torch.zeros_like(prob)
p = torch.where(prob < 0.3, t, f)
print(p.requires_grad) # prints False

However if you set requires_grad to True fo…