transformer_engine.pytorch.distributed.checkpoint function only works with TE modules, instead of all Callables

[MaciejBalaNV]

The activaction checkpointing function only works for TE modules, even though the argument function is typed as Callable, indicating it should work with an arbitrary function, like PyTorch activation checkpointing does. The error in TE comes from this line, where it's assumed the function argument needs fsdp_wrapped attribute. The error message when we pass a function, or a method in this case, is as follows:

[rank0]:   File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/distributed.py", line 651, in checkpoint
[rank0]:     setattr(function, "fsdp_wrapped", False)
[rank0]: AttributeError: 'method' object has no attribute 'fsdp_wrapped'