class MultiScaleEdgeEnergy(torch.nn.Module):
def __init__(self):
- super(MultiScaleEdgeEnergy, self).__init__()
+ super().__init__()
k = torch.exp(- torch.tensor([[-2., -1., 0., 1., 2.]])**2 / 2)
k = (k.t() @ k).view(1, 1, 5, 5)
self.gaussian_5x5 = torch.nn.Parameter(k / k.sum()).requires_grad_(False)