Hello @albanD, I have converted the custom autograd.Function code as new way, but I can’t save the torch model, do you have any suggests pls ?
class bin(torch.autograd.Function):
    """bin {0, 1} a real valued tensor."""

    @staticmethod
    def forward(ctx, inputs, threshold=DEFAULT_THRESHOLD):
        ctx.threshold = threshold
        outputs = inputs.clone()
        outputs[inputs.le(ctx.threshold)] = 0
        outputs[inputs.gt(ctx.threshold)] = 1
        return outputs

    @staticmethod
    def backward(ctx, gradOutput):
        return gradOutput, None


torch.save(ckpt, savename)
File “/home/quannguyen/.local/lib/python3.6/site-packages/torch/serialization.py”, line 372, in save
_save(obj, opened_zipfile, pickle_module, pickle_protocol)
File “/home/quannguyen/.local/lib/python3.6/site-packages/torch/serialization.py”, line 476, in _save
pickler.dump(obj)
TypeError: can’t pickle bin objects