diff --git a/modules/transformation.py b/modules/transformation.py index c43f173..311838a 100755 --- a/modules/transformation.py +++ b/modules/transformation.py @@ -4,6 +4,7 @@ import torch.nn as nn import torch.nn.functional as F device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + class TPS_SpatialTransformerNetwork(nn.Module): """ Rectification Network of RARE, namely TPS based STN """