TracerWarning error when calling SummaryWriter add_graph

When I call SummaryWriter().add_graph(model, input), I get the following error:

TracerWarning: Output nr 1. of the traced function does not match the corresponding output of the Python function. Detailed error:
With rtol=1e-05 and atol=1e-05, found 32 element(s) (out of 32) whose difference(s) exceeded the margin of error (including 0 nan comparisons). The greatest difference was 0.013833858072757721 (-0.04951992258429527 vs. -0.03568606451153755), which occurred at index (1, 0).

I Initially suspected that the dropout layers may be the cause, but setting model.train(False) before calling add_graph doesn’t resolve the issue.

Pytorch 1.9.1
TensorBoard 2.7.0