LSTM incompatible with vmap? "accessing '.data' under vmap transform is not allowed"

Been trying to figure out this problem for quite some time now and it’s proving quite intractable, and searches on the error return no results…apparently a very rare problem, as it seems no one anywhere has ever posted anything about this particular error message (at least that I’ve been able to find).

Basically the problem is that a particular step in executing the forward pass through an LSTM module seems to be fundamentally incompatible with Torch’s vmap apparatus.

A full description of the problem with a minimal reproducible case and full stack trace can be found here: vmap and rnn/lstm "accessing '.data' under vmap transform is not allowed" · Issue #105982 · pytorch/pytorch · GitHub

Any takers? This problem has proven intractable for me for months at this point.