Solved by Mona_Jalal in post #3 
              
              
                changed line 118 of type_specific_network.py 
from: 
norm = torch.norm(masked_embedding, p=2, dim=2) + 1e-10 
to: 
norm = torch.norm(masked_embedding, p=2, dim=2, keepdim=True) + 1e-10