Backtracking my answer on the Github issue here for anyone who might be facing the same kind of issue:
The problem arises in modeling_openai.py when the user do not provide the  position_ids  function argument thus leading to the inner  position_ids  being created during the forward call. This is fine in classic PyTorch because  forward  is actually evaluated at each call. When it comes to tracing, this is an issue, because the device specified in the forward to actually create the tensor will be hardcoded and you can actually see it in the generated graph:
%input.1 : Tensor = aten::view(%input_ids.1, %64) 
%140 : Device = prim::Constant[value="cuda:0"]() 
%position_ids.1 : Tensor = aten::arange(%59, %67, %45, %140, %70) 
%73 : Tensor = aten::unsqueeze(%position_ids.1, %45)

Above you can see  %140  is a constant which value is actually set to  "cuda:0"  and then, it is reused to create the  %position_ids.1  tensor through  aten::arange(..., %140, ...)  which of course leads to the error you’re seeing.
I’ll have a fix to generate the  position_ids  buffer correctly registered at the Module initialisation and not during forward, so it should be correctly handled by the  map_location  parameter while exporting.