issue with checkpoint saving

[mijanur132]

The error arise from:
transfusion.py:779----> if learned_value_residual_mix else (lambda _: 0.5)

Error:

rank0: File "/ccs/home/palashmr/packages/miniconda/pyt_env/transfusion/lib/python3.12/site-packages/torch/serialization.py", line 780, in save
rank0: _save(
rank0: File "/ccs/home/palashmr/packages/miniconda/pyt_env/transfusion/lib/python3.12/site-packages/torch/serialization.py", line 1004, in _save
rank0: pickler.dump(obj)
rank0: AttributeError: Can't get local object 'Attention.init.locals.lambda"

Potential solution:
to replace lambda function with a regular function.

[lucidrains]

@mijanur132 ah, how about after upgrading to 0.6.5?