So I worked on making a minimal example of this, and in fiddling with the parameters I found that reducing the minibatch size makes the error go away. I’m glad it’s fixed but this doesn’t make sense to me - should changing the minibatch size change my computation graph? Based on the error message I would never have guessed it was related to my batch size.
The simplified version of the code I made is here 2, though the fix is the same as in the original.