Yes, I do use indexing. I managed to strip down my module code, unrolled a few loops and ended up with this snippet:
import torch
import torch.nn as nn


class Core(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.f = nn.Linear(4, 2)

    def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        x = torch.cat((x1, x2), 1)
        return self.f(x)


class MyFancyModel(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.core = Core()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        assert x.shape == (1, 2, 2)

        y = torch.empty(1, 4, 2)
        y[:, :2] = x

        x1, x2 = x[:, 0], x[:, 1]
        y1, y2 = y[:, 0], y[:, 1]
        y[:, 2] = x2 + (x2 - x1) * self.core(y1, y2)

        x1, x2 = x[:, 1], y[:, 2]
        y1, y2 = y[:, 1], y[:, 2]
        y[:, 3] = x2 + (x2 - x1) * self.core(y1, y2)

        return y

In case my indexing is broken, can you give me a hint how to fix it?