Megatron-DeepSpeed

87a9dba0 - `torch.testing.assert_equal` didn't make it (#273)

`torch.testing.assert_equal` didn't make it (#273) looks like pt-1.11 dropped `torch.testing.assert_equal`, so using `torch.testing.assert_equal` instead