Problem with PT_MAP

[callumHub]

Describe the bug
Using PT-MAP with code adapted from episodic training.ipynb, recieve error message:
RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor [150, 5]], which is output 0 of ExpBackward0, is at version 579; expected version 0 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True).
To Reproduce
Steps to reproduce the behavior:

1. Python version 3.12
2. Episodic training
3. Features Dataset

Additional context
'Fixed' error message by changing all in place operations in pt_map.py. EX: x *= y -> x = x*y.

I've tried most other FSL algorithms here and not experienced this problem. Seems to be in the compute_optimal_transport method.

[callumHub]

Also: Using MSELoss for PT_MAP does not seem to work. I'm pretty sure I have formated the labels and model outputs correctly, just the MSELoss does not go down during training. Using CELoss allows for training, but that is not what the paper had said, and the accuracy is low compared to the other FSL algorithms.

Furthermore, with MSELoss, the sinkhorn algorithm is very unstable, the transport plans values get very small and cause loss to become nan, If I clip the transport plan, all values become the same => model can not learn. I've tried different settings of lambda_regularization but the loss either stays at its original value or becomes nan.