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