Misc updates to transforms, loss and UNETR

[constantinpape]

No description provided.

[constantinpape]

This contains a couple recent changes, the most important thing is an update to allow passing a state dict to UNETR as well. This makes it easier to change encoder checkpoints, so that they don't have to be resaved and so that we don't need to update the code to parse the checkpoints every time.

@anwai98 with these changes loading the checkpoint from @cdalinghaus still fails with the following error:

RuntimeError: Error(s) in loading state_dict for ViT_MAE:                                                                                                                              
        Missing key(s) in state_dict: "head.weight", "head.bias".                                                                                                                      
        size mismatch for pos_embed: copying a param with shape torch.Size([1, 197, 1024]) from checkpoint, the shape in current model is torch.Size([1, 4097, 1024]).