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([])