I changed the function into this:
    def relaxed_taylor_closed_form_solution(ctx, w):
        w_shift_left = torch.cat((w[:, 1:], w[:, -1].reshape(-1, 1)), 1)
        w_shift_right = torch.cat((w[:, 0].reshape(-1, 1), w[:, :-1]), 1)
        w_shift_up = torch.cat((w[1:, :], w[-1, :].reshape(1, -1)), 0)
        w_shift_down = torch.cat((w[0, :].reshape(1, -1), w[:-1, :]), 0)
        interactions = ctx.alpha*(torch.sign(w - w_shift_left) + torch.sign(w - w_shift_right) + torch.sign(w - w_shift_up) + torch.sign(w - w_shift_down))
        return ctx.beta - interactions

The execution was still too slow, so I started investigating on parallelization. Currently, I implemented this, which works good for CPU but crashes with GPU:
from joblib import Parallel, delayed
import multiprocessing

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

    def forward(self, x):
        batch_size, channels, height, width = x.shape
        num_cores = multiprocessing.cpu_count()
        results = Parallel(n_jobs=num_cores, backend="threading")(delayed(FrankWolfe_MinCutSum_Canonical.apply)(x[b, c]) for b in range(batch_size) for c in range(channels))
        out = torch.stack(results).reshape(batch_size, channels, height, width)
        return out

To give you a general idea, the execution is faster for parallelized CPU than for GPU without parallelization. It would be ideal to make that part of the code to be called in parallel, since the calls are completely independent and only need to be synchronized before passing to the next layer.