This shouldn’t make any difference as seen here:
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = nn.Linear(10, 10, bias=False)
        self.fc2 = nn.Linear(10, 3, bias=False)
        
    def forward(self, x, return_slice=False):
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        if return_slice:
            return x[:, -1]
        else:
            return x

model = MyModel()
x = torch.randn(10, 10)

# Complete batch approach
out = model(x)
out[:, -1].mean().backward()

for name, param in model.named_parameters():
    print(name, param.grad.abs().sum())

# Sliced approach
model.zero_grad()
out = model(x, return_slice=True)
out.mean().backward()

for name, param in model.named_parameters():
    print(name, param.grad.abs().sum())

You will get the same gradients for both methods.
How reproducible is this issue using your code?
I.e. are you always getting the good results in the former approach and bad ones in the latter?