for stemgnn, does

n_series

always have to match

len(df['unique_id'].distinct())

?

Hey, yes. All series must be present in every batch

yah okay makes sense. we're finding the memory requirements around ~500 series with stemgnn to be quite high even with low to moderate input sizes

I think all the possible windows are generated upfront during training and then just sampled, this is something that we have to fix. Can you please open an issue so we can track it?