Both approaches should yield the same result and I assume your models might have used different layer parameters.
This code snippet shows that the different forward implementations create the same gradients up the floating point precision:
import torch
torch.manual_seed(0)
import torch.nn as nn
import torch.nn.functional as F

class Flatten(nn.Module):
    def forward(self, input):
        return input.view(input.size(0), -1)

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.pre_heads = nn.ModuleDict()
        for head_name in ['a','b','c']:
                self.pre_heads[head_name] = nn.Sequential(
                    nn.Conv2d(16,16, kernel_size=1),
                    nn.ReLU()
                )
        self.out = nn.Sequential(
            Flatten(),
            nn.Linear(16*4*4,1)
        )
    
    def forward(self, x, head_name, use_heads=True):
        if use_heads:
            x = torch.cat(
                [self.pre_heads[j](x[i].unsqueeze(0)) for i, j in enumerate(head_name)],
                dim=0
            )
            x = torch.squeeze(torch.sigmoid(self.out(x)))
            return x
        else:
            x = self.pre_heads['a'](x)
            x = torch.squeeze(torch.sigmoid(self.out(x)))
            return x

mymodel = Model()
sample_input = torch.ones(4,16,4,4)

# 
output = mymodel(sample_input, head_name=['a', 'a', 'a', 'a'])
loss = F.binary_cross_entropy(output, torch.Tensor([1,1,1,1]))
loss.backward()
grad0 = mymodel.pre_heads['a'][0].weight.grad.sum().clone()

#
mymodel.zero_grad()
output = mymodel(sample_input, head_name=None, use_heads=False)
loss = F.binary_cross_entropy(output, torch.Tensor([1,1,1,1]))
loss.backward()
grad1 = mymodel.pre_heads['a'][0].weight.grad.sum().clone()

print((grad0 - grad1).abs())
> tensor(1.1921e-07)