#26 fix training and inference

[afmjoaa]

• [x] Solve multiple backward calls for output classification heads
• [x] Add Weight in loss functions
• [x] Use ADAMW instead of ADAM as optimizer
• [x] Solved the saved model loading issue