Update training logic to computes loss over logits#990

anwai98 wants to merge 7 commits intomastercomputational-cell-analytics/micro-sam:masterfrom

train-on-low-rescomputational-cell-analytics/micro-sam:train-on-low-resCopy head branch name to clipboard