Update modelling_magi.py

modelling_magi.py

@@ -449,7 +449,7 @@ class MagiModel(PreTrainedModel):
             affinity_matrices = []
             for crop_embeddings in crop_embeddings_for_batch:
                 crop_embeddings = crop_embeddings / crop_embeddings.norm(dim=-1, keepdim=True)
-                affinity_matrix = einsum("i d, j d -> i j", affinity_matrix)
+                affinity_matrix = crop_embeddings @ crop_embeddings.T
                 affinity_matrices.append(affinity_matrix)
             return affinity_matrices
         affinity_matrices = []