Adapting self-supervision tutorial for 2D case - contrastive loss

[ptvu281]

I'm trying to adapt the self-supervision tutorial for a 2D case. However, I keep getting a tensor shape mismatch error at the contrastive loss step. I've posted a screenshot of the error here: https://imgur.com/a/AXmWKgB
Looking at the source code, I see that the error could be caused by the shape mismatch between negatives_mask and sim_matrix but I'm not sure how to proceed. The things I have changed to adapt for 2D are:

• At the flatten step during training, change end_dim=4 to end_dim=3
• Change img_size and patch_size to a 2-tuple instead of 3-tuple

Any suggestion on how to proceed is appreciated!!