Simplest policy gradient code implementation assumes uniform episode lengths?

The simplest policy gradient formula is given by

\hat{g} = \frac{1}{|\mathcal{D}|} \sum_{\tau \in \mathcal{D}} \sum_{t=0}^{T} \nabla_{\theta} \log \pi_{\theta}(a_t |s_t) R(\tau),

In the relevant part of the code, it is implemented as:

def compute_loss(obs, act, weights):
    logp = get_policy(obs).log_prob(act)
    return -(logp * weights).mean()

In this code, the dimensions of weights and logp is batch_size*(sum of all episode lengths). This loss will only be proportional to policy gradient if all the episode lengths are the same, which is very unlikely. Typically in any given batch, different episodes will have different lengths and and we should sum the logp*weights for each episode separately before taking the batch mean?

Am I missing something?