def METHOD_NAME(self, loss_type, loss_args, forward_args):
    """
    The goal of this test is to assess if the gradient of the loss function
    is correct by testing if we can train a one layer neural network
    to segment one image.
    We verify that the loss is decreasing in almost all SGD steps.
    """
    learning_rate = 0.001
    max_iter = 40
    # define a simple 3d example
    target_seg = torch.tensor(
        [
            [
                # raw 0
                [[0, 0, 0, 0], [0, 1, 1, 0], [0, 1, 1, 0], [0, 0, 0, 0]],
                # raw 1
                [[0, 0, 0, 0], [0, 1, 1, 0], [0, 1, 1, 0], [0, 0, 0, 0]],
                # raw 2
                [[0, 0, 0, 0], [0, 1, 1, 0], [0, 1, 1, 0], [0, 0, 0, 0]],
            ]
        ],
        device=self.device,
    )
    target_seg = torch.unsqueeze(target_seg, dim=0)
    image = 12 * target_seg + 27
    image = image.float().to(self.device)
    num_classes = 2
    num_voxels = 3 * 4 * 4
    target_onehot = one_hot(target_seg, num_classes=num_classes)
    # define a one layer model
    class OnelayerNet(nn.Module):
        def __init__(self):
            super().__init__()
            self.layer_1 = nn.Linear(num_voxels, 200)
            self.acti = nn.ReLU()
            self.layer_2 = nn.Linear(200, num_voxels * num_classes)
        def forward(self, x):
            x = x.view(-1, num_voxels)
            x = self.layer_1(x)
            x = self.acti(x)
            x = self.layer_2(x)
            x = x.view(-1, num_classes, 3, 4, 4)
            return x
    # initialise the network
    net = OnelayerNet().to(self.device)
    # initialize the loss
    loss = loss_type(**loss_args)
    # initialize a SGD optimizer
    optimizer = optim.Adam(net.parameters(), lr=learning_rate)
    loss_history = []
    init_output = None
    # train the network
    for iter_i in range(max_iter):
        # set the gradient to zero
        optimizer.zero_grad()
        # forward pass
        output = net(image)
        if init_output is None:
            init_output = torch.argmax(output, 1).detach().cpu().numpy()
        if loss_args["to_onehot_y"] is False:
            loss_val = loss(output, target_onehot, **forward_args)
        else:
            loss_val = loss(output, target_seg, **forward_args)
        if iter_i % 10 == 0:
            pred = torch.argmax(output, 1).detach().cpu().numpy()
            gt = target_seg.detach().cpu().numpy()[:, 0]
            print(f"{loss_type.__name__} iter: {iter_i}, acc: {np.sum(pred == gt) / np.prod(pred.shape)}")
        # backward pass
        loss_val.backward()
        optimizer.step()
        # stats
        loss_history.append(loss_val.item())
    pred = torch.argmax(output, 1).detach().cpu().numpy()
    target = target_seg.detach().cpu().numpy()[:, 0]
    # initial predictions are bad
    self.assertTrue(not np.allclose(init_output, target))
    # final predictions are good
    np.testing.assert_allclose(pred, target)