The values of squared_norm in PrimaryCaps explode and create the NaNs.
In the last iteration before the NaNs are raised PrimaryCaps creates tensors with these statistics:
print(squared_norm.min(), squared_norm.max())
> tensor(1.4527e+20, device='cuda:0', grad_fn=<MinBackward1>) tensor(5.3600e+26, device='cuda:0', grad_fn=<MaxBackward1>)
print(input_tensor.min(), input_tensor.max())
> tensor(-1.6448e+13, device='cuda:0', grad_fn=<MinBackward1>) tensor(1.8647e+13, device='cuda:0', grad_fn=<MaxBackward1>)
print(output_tensor.min(), output_tensor.max())
> tensor(nan, device='cuda:0', grad_fn=<MinBackward1>) tensor(nan, device='cuda:0', grad_fn=<MaxBackward1>)

This can be explained via:
squared_norm = torch.tensor(1e26)
input_tensor = torch.tensor(1e13)
output_tensor = squared_norm * input_tensor / ((1+squared_norm) * torch.sqrt(squared_norm)) # divisor is Inf
print(output_tensor)
> tensor(nan)

So it seems your overall training might not be stable.