You need to fix the value of output_dim beforehand so you can make use of the same output weight matrix always (and so that weight updates are informative). The way I would do this is to make a guess at the largest output size you will need (max of training data, 1.5 times that, …?) and then mask the output for smaller target values. This means that you multiply the last k scalars of your output by 0, the ones which you dont need. The specifics of when to mask depend on when you know about your current output size. Simple example: If you already know you have a target tensor of a specific size beforehand, you could do:
max_size = big
for x, y in batches:
    mask_last_k = max_size-y.shape[-1]
    mask = torch.ones(max_size)
    mask[mask_last_k:] = 0
    
    predictions = model(x)
    predictions_masked = predictions * mask
    
    loss = loss_fn(predictions_masked, y)
    loss.backward()
    optimizer.step()

If you actually somehow generate output tensors of different sizes, you can of course also pad the difference up to the required length:
    pad_last_k = y.shape[-1] - x.shape[-1]
    pad = torch.zeros(max_size)
    
    predictions = model(x)
    predictions_padded = torch.cat([predictions, mask], dim=-1)

If you get autograd errors for creating the pad/ mask tensor yourself, create them in the scope of with torch.no_grad().
Hope this helps