Update models.py
Browse files
models.py
CHANGED
|
@@ -12,6 +12,7 @@ import torch
|
|
| 12 |
from einops import rearrange
|
| 13 |
from scipy import ndimage
|
| 14 |
from torch.special import gammaln
|
|
|
|
| 15 |
|
| 16 |
class AlignmentEncoder(torch.nn.Module):
|
| 17 |
"""
|
|
|
|
| 12 |
from einops import rearrange
|
| 13 |
from scipy import ndimage
|
| 14 |
from torch.special import gammaln
|
| 15 |
+
import torch.nn as nn
|
| 16 |
|
| 17 |
class AlignmentEncoder(torch.nn.Module):
|
| 18 |
"""
|