Upload hierarchy_model.py with huggingface_hub

hierarchy_model.py

@@ -282,7 +282,19 @@ class HierarchyEncoder(nn.Module):
     
     def forward(self, hierarchy_indices):
         # hierarchy_indices: (B,) - batch of hierarchy indices
-        emb = self.embedding(hierarchy_indices)
+        # Workaround for MPS: embedding layers don't work well with MPS, so do lookup on CPU
+        device = next(self.parameters()).device
+        if device.type == 'mps':
+            # Move indices to CPU for embedding lookup
+            indices_cpu = hierarchy_indices.cpu()
+            # Use functional embedding with explicit weight handling for MPS compatibility
+            emb_weight = self.embedding.weight.cpu()
+            emb = F.embedding(indices_cpu, emb_weight)
+            # Move result back to model device (MPS) - ensure it's contiguous
+            emb = emb.contiguous().to(device)
+        else:
+            emb = self.embedding(hierarchy_indices)
+        # Ensure emb is on the same device as projection before calling it
         return self.projection(emb)
 
 class HierarchyClassifierHead(nn.Module):