Undefined Gradients due to division by zero.

[saswatpp]

https://github.com/jindongwang/Pytorch-CapsuleNet/blob/93996119a224e9c3d274df88735f1acb36592d97/capsnet.py#L79

To make it stable, following is suggested:

def squash(self, input_tensor, epsilon=1e-7):
    squared_norm = (input_tensor ** 2 + epsilon).sum(-1, keepdim=True)
    output_tensor = squared_norm * input_tensor / ((1. + squared_norm) * torch.sqrt(squared_norm))
    return output_tensor