//  [python]
def test_groupnorm_nhwc(self):
        def helper(self, size, groups, memory_format, dtype, is_mixed):
            channels = size[1]
            input = torch.randn(size, dtype=dtype, requires_grad=True)
            input = input.contiguous(memory_format=memory_format)
            input.retain_grad()
            grad = torch.randn(size, dtype=dtype)
            grad = grad.contiguous(memory_format=memory_format)
            if dtype == torch.bfloat16 and is_mixed:
                gn = nn.GroupNorm(groups, channels).to(torch.float)
            else:
                gn = nn.GroupNorm(groups, channels).to(dtype)
            gn.weight.data.uniform_()
            gn.bias.data.uniform_()

            ref_input = (
                input.detach()
                .clone()
                .contiguous(memory_format=torch.contiguous_format)
                .requires_grad_(True)
            )
            ref_grad = (
                grad.detach().clone().contiguous(memory_format=torch.contiguous_format)
            )
            if dtype == torch.bfloat16 and is_mixed:
                ref_gn = nn.GroupNorm(groups, channels).to(torch.float)
            else:
                ref_gn = nn.GroupNorm(groups, channels).to(dtype)
            ref_gn.load_state_dict(gn.state_dict())

            out = gn(input)
            out.backward(grad)
            ref_out = ref_gn(ref_input)
            ref_out.backward(ref_grad)

            self.assertTrue(out.is_contiguous(memory_format=memory_format))
            self.assertTrue(
                ref_out.is_contiguous(memory_format=torch.contiguous_format)
            )
            torch.testing.assert_close(out, ref_out)
            # parameters in bfloat16/Half is not recommended
            atol = 5e-4
            rtol = 8e-3
            torch.testing.assert_close(
                gn.weight.grad, ref_gn.weight.grad, atol=atol, rtol=rtol
            )
            torch.testing.assert_close(
                gn.bias.grad, ref_gn.bias.grad, atol=atol, rtol=rtol
            )
            torch.testing.assert_close(input.grad, ref_input.grad, atol=atol, rtol=rtol)

        for dtype in [torch.float16, torch.bfloat16, torch.float, torch.double]:
            for is_mixed in [True, False]:
                helper(self, (4, 8, 10, 10), 4, torch.channels_last, dtype, is_mixed)
                helper(self, (2, 30, 9, 9), 3, torch.channels_last, dtype, is_mixed)
                helper(self, (4, 8, 40, 40), 4, torch.channels_last, dtype, is_mixed)
                helper(self, (4, 40, 40, 40), 2, torch.channels_last, dtype, is_mixed)
                helper(self, (2, 30, 50, 50), 3, torch.channels_last, dtype, is_mixed)
                helper(self, (2, 60, 50, 50), 3, torch.channels_last, dtype, is_mixed)
                helper(
                    self, (2, 9, 7, 11, 15), 3, torch.channels_last_3d, dtype, is_mixed
                )
                helper(
                    self, (2, 9, 7, 200, 15), 3, torch.channels_last_3d, dtype, is_mixed
                )
                helper(
                    self,
                    (2, 60, 7, 200, 15),
                    3,
                    torch.channels_last_3d,
                    dtype,
                    is_mixed,
                )