Hi @tom, I tried implementing what you have mentioned.
threshold=nn.Parameter(torch.tensor(1.5))
threshold.requires_grad=True
eps=1e-7
z=torch.zeros((),dtype=a.dtype)
a_new=torch.where((a-threshold).abs()<eps,a,z)

Now I try to do the backward pass
a_new.sum().backward()

But I am getting RuntimeError:
element 0 of tensors does not require grad and does not have a grad_fn