The 'use_softmax' arg is missing in FocalLoss in the DiceFocalLoss class. See the screenshot. <img width="556" alt="Screenshot 2024-11-28 at 10 42 24" src="https://github.com/user-attachments/assets/742b2883-9d8d-4845-acc9-92bc4c3077b4">