Update models.py

models.py

@@ -12,6 +12,7 @@ import torch
 from einops import rearrange
 from scipy import ndimage
 from torch.special import gammaln
+import torch.nn as nn
 
 class AlignmentEncoder(torch.nn.Module):
     """