@Juans Thanks for clarifying your use case.
It sounds like multiprocessing.Pool.map (or something similar) would allow you to map a function like the f_i in your example on some chunk of the input tensor to a process pool that you can define. The exact chunksize can be configured using args like chunksize or a related operation called imap. Will this serve your purpose?
I tried the following (very simplistic) example using Pool.map in the model’s forward pass and it seems to work:
import torch
import torch.nn as nn
import torch.optim as optim
import torch.multiprocessing as mp

class DemoMultiProcModel(nn.Module):
    def __init__(self):
        super(DemoMultiProcModel, self).__init__()
        self.net1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(10, 5)
        
    def forward(self, x):    
        with mp.Pool(5) as p:
            self.result = p.map(sum, x)
        return self.net2(self.relu(self.net1(x)))


def train():
    print(f"Running Model with parallelized forward pass.")
    model = DemoMultiProcModel()

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(model.parameters(), lr=0.001)
    
    for batch_idx in range(5):
        # Generate random inputs/labels
        inputs = torch.randn(20, 10)
        labels = torch.randn(20, 5)
        
        # Train
        optimizer.zero_grad()
        outputs = model(inputs)
        loss_fn(outputs, labels).backward()
        optimizer.step()
        print(f"Batch {batch_idx} done.")