Use torch.all to check tensor equality

Use torch.all to check tensor equality #621