Is this the correct way?
    def relaxed_taylor_closed_form_solution(ctx, w):
        u_star = ctx.beta
        u_star -= (torch.cuda if torch.cuda.is_available() else torch).FloatTensor([ \
                    ctx.alpha*sum([(1 if w[(x, y)] > w[j] else -1) for j in FrankWolfe_MinCutSum_Canonical.neighbors(ctx, (x, y))]) \
                    for x in range(ctx.width) for y in range(ctx.height)]).reshape(ctx.width, ctx.height)
        return u_star

class LateralInteractions(nn.Module):
    def __init__(self):
        super(LateralInteractions, self).__init__()

    def forward(self, x):
        # out = x
        batch_size, channels, height, width = list(x.shape)
        out = torch.stack([FrankWolfe_MinCutSum_Canonical.apply(Variable(x[b, c], requires_grad=True)) \
                for b in range(batch_size) for c in range(channels)]).reshape(batch_size, channels, height, width)
        # out[b, c] = FrankWolfe_MinCutSum_Canonical.apply(x[b, c])
        return out

Everything else is as in the original post.