As your function’s derivative is 0 everywhere (except for the input 0, where it isn’t smooth), you can also implement it as a function:
class StepFunction(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input):
        return (input > 0).float()
    
    @staticmethod
    def backward(ctx, grad_output):
        return torch.zeros_like(grad_output)

w = torch.FloatTensor([1.0, 2.0])
w.requires_grad = True
step_function = StepFunction.apply
out = w * 2
out = step_function(out)

out.sum().backward()