Thank you for the hints. I tried to do the BMM implementation because the inputs are dense continuous values. I think I got close, but…
class ModelParallel2(nn.Module):
    def __init__(self, in_features, hidden, out_features):
        super().__init__()
        
        self.out_features = out_features
        
        
        self.block1_bias = nn.Parameter(torch.ones(hidden))
        self.block1 = nn.Parameter(torch.ones(
            1, in_features, hidden
        ))
        
        self.block2_bias = nn.Parameter(torch.ones(hidden))
        self.block2 = nn.Parameter(torch.ones(
            1, hidden, hidden
        ))
        
        self.block3_bias = nn.Parameter(torch.ones(1))
        self.block3 = nn.Parameter(torch.ones(
            1, hidden, 1
        ))
        
        
        # nn.init.xavier_uniform_(self.block1, gain=nn.init.calculate_gain('relu'))
        # nn.init.xavier_uniform_(self.block2, gain=nn.init.calculate_gain('relu'))
        # nn.init.xavier_uniform_(self.block3, gain=nn.init.calculate_gain('relu'))

        nn.init.kaiming_uniform_(self.block1, a=np.sqrt(6))
        nn.init.kaiming_uniform_(self.block2, a=np.sqrt(6))
        nn.init.kaiming_uniform_(self.block3, a=np.sqrt(6))
        
        bound_a = 1 / np.sqrt(in_features)
        bound_b = 1 / np.sqrt(hidden)
        nn.init.uniform_(self.block1_bias, -bound_a, bound_a)
        nn.init.uniform_(self.block2_bias, -bound_b, bound_b)
        nn.init.uniform_(self.block3_bias, -bound_b, bound_b)
        
    def forward(self, x):
        # torch.Size([32, 130])
        bs = x.shape[0]
        
        # torch.Size([32, 1, 130])
        x = x[:, None, :]

        # torch.Size([32, 45, 130])
        x = x.expand(bs, self.out_features, x.shape[2])

        # torch.Size([32, 45, 10])
        x = x @ self.block1
        x = F.relu(x, inplace=True)
        
        # torch.Size([32, 45, 10])
        x = x @ self.block2
        x = F.relu(x, inplace=True)
        
        # torch.Size([32, 45, 1])
        x = x @ self.block3

        # torch.Size([32, 45])
        x = x.squeeze(dim=-1)
        
        return x

Inputting a random tensor and initializing the model with:
mp = ModelParallel2(
    in_features=130,
    hidden=10,
    out_features=45
)

I get a torch.Size([32, 45]) sized output, which is correct and desirable. However, the output features are all duplicated:
tensor([[-0.0014, -0.0014, -0.0014,  ..., -0.0014, -0.0014, -0.0014],
        [ 0.0002,  0.0002,  0.0002,  ...,  0.0002,  0.0002,  0.0002],
        [ 0.0001,  0.0001,  0.0001,  ...,  0.0001,  0.0001,  0.0001],
        ...,
        [-0.0023, -0.0023, -0.0023,  ..., -0.0023, -0.0023, -0.0023],
        [-0.0003, -0.0003, -0.0003,  ..., -0.0003, -0.0003, -0.0003],
        [-0.0006, -0.0006, -0.0006,  ..., -0.0006, -0.0006, -0.0006]],
       grad_fn=<SqueezeBackward1>)

From the docs I found out that torch.bmm “This function does not broadcast. For broadcasting matrix products, see torch.matmul().” But even replacing x = x @ self.block1 with x = torch.matmul(x, self.block1) doesn’t change the behavior of the duplicates.