as_torch_callback function

Convert to a TorchCallback