Thanks for the input. I rewrote the squash function as given below
def squash(self, input_tensor):
        squared_norm = (input_tensor ** 2).sum(-1, keepdim=True)
        denom = ((1. + squared_norm) * torch.sqrt(squared_norm))
        if torch.isinf(denom).sum().item()>0:
          output_tensor = input_tensor / torch.sqrt(squared_norm)
        else:
          output_tensor = squared_norm * input_tensor / ((1. + squared_norm) * torch.sqrt(squared_norm))
        return output_tensor

I think this solves it.