3

I traced my Neural Network using torch.jit.trace on a CUDA-compatible GPU server. When I reloaded that Trace on the same server, I could reload it and use it fine. Now, when I downloaded it onto my laptop (for quick testing), when I try to load the trace I get:

RuntimeError: Could not run 'aten::empty_strided' with arguments from the 'CUDA' backend. 'aten::empty_strided' is only available for these backends: [CPU, BackendSelect, Named, AutogradOther, AutogradCPU, AutogradCUDA, AutogradXLA, AutogradPrivateUse1, AutogradPrivateUse2, AutogradPrivateUse3, Tracer, Autocast, Batched, VmapMode].

Can I not switch between GPU and CPU on a trace? Or is there something else going on?

Mike B
  • 2,136
  • 2
  • 12
  • 31
JCunn
  • 53
  • 1
  • 7

1 Answers1

2

I had this exact same issue. In my model I had one line of code that was causing this:

if torch.cuda.is_available():
    weight = weight.cuda()

If you have a look at the official documentation for trace (https://pytorch.org/docs/stable/generated/torch.jit.trace.html) you will see that

the returned ScriptModule will always run the same traced graph on any input. This has some important implications when your module is expected to run different sets of operations, depending on the input and/or the module state

So, if the model was traced on a machine with GPU this operation will be recorded and you won't be able to even load your model on a CPU only machine. To solve this, deleted everything that makes you model CUDA dependent. In my case it was as easy as deleting the code-block above.

Mike B
  • 2,136
  • 2
  • 12
  • 31
  • 1
    Yeah, this gets a little tricky if you have to create tensors within your module(like random tensors) instead of just processing, but generally this is good advice – Juliano S Assine Jan 13 '23 at 18:49