Eventually solved the problem, just leaving this here in case anyone else is ever in the same boat. When I traced my model in pytorch I was sending my model to the GPU. I changed this to the CPU and everything started working. I don’t know why, I don’t know if there was something else going on, all I know is that things started working after I made this change.