Hi!
First - thank you for your responses. @JuanFMontesinos - that is just a mws - in reality, the code has different computation steps for different samples (so some form of conditional computation model). Now that I know that parallelizing that loop is not possible (@albanD - thank you here, I understand your argument) , I’ve instead settled for the following solution (in case anybody else ever stumbles across a similar problem):
class Model(nn.Module):
    def __init__(self):
        nn.Module.__init__(self)
        self.lin1 = nn.Linear(100, 100)
        self.lin2 = nn.Linear(100, 30)
        self.lin3 = nn.Linear(30, 3)

    def forward(self, xs):
        step1 = self.lin1(xs)
        step2 = []
        results = torch.zeros(*results_shape).to(xs.device)
        # computing different masks for different computational steps
        mask1 = some_masking_function1(xs)
        results[mask] = some_computation1(xs[mask])
        mask2 = some_masking_function2(xs)
        results[mask] = some_computation2(xs[mask])
        ...
        maskn = some_masking_functionn(xs)
        results[mask] = some_computationn(xs[mask])
        results = self.lin3(results)
        return results

This solution effectively “clusters” samples by which computation I want to perform, and then batches over all these samples. As I have potentially many computation steps, this approach did only help me by a factor of 1.5-1.8 - but I’ll take it!
Thank you again.