def test_l2_gpu():
    prediction = torch.rand(*test_shape).to('cuda:0')
    loss = F.contextual_loss(prediction, prediction, loss_type='l2')
    assert loss.shape == torch.Size([])
def test_l1():
    prediction = torch.rand(*test_shape)
    loss = F.contextual_loss(prediction, prediction, loss_type='l1')
    assert loss.shape == torch.Size([])