cudatoolkit               10.2.89
torch                     1.9.0

These are my CUDA and torch versions, I am certain PyTorch is current since I upgraded recently, within the last 48h, unsure about CUDA but I think so.

On the first point, you may have been right. After doing what I thought was the fix, I still ran into this error:

Traceback (most recent call last):
  File "/panfs/roc/groups/13/suo-yang/dikem003/DimensionReductionNLE/auto_ode/AETrainingConditionNum.py", line 238, in <module>
    stiffness_loss = AE_stiffness_loss(latent_predictions) * EPOCH_SCALER
RuntimeError: CUDA error: an illegal memory access was encountered

After it failed again, I added some debugging statements in my main training loop to see the sizes of returned matrices from both loss functions (while changing nothing else):
# print shapes
print('\nRL Size',reconstruction_loss.size())
print('\nSL Size',stiffness_loss.size())

Now suddenly the error does not appear, as if referencing the size of both loss tensors before I call backward has eliminated the issue altogether. I even checked my git diff to see if I changed anything else and I have not, I only added print statements. I have made it to ~100 epochs at the time of writing whereas it was failing consistently at around 5-10 before.
I can try to develop a minimal reproducible program with this loss function since I know my one example isn’t worth much, I will need to figure out how to trigger the noninvertibility exception without my specialized set of data. Right now I sort of feel like I’m chasing Schrodinger’s bug, driving me nuts. I don’t see any reason why checking the loss tensor size in my training loop would affect anything.

The shape issue for torch.cond is reflected in the documentation, and I can also generate a small program to demonstrate this and raise an Github issue, no problem there.