# Loaded Torchscript models have wrong output shape
**URL:** https://discuss.ray.io/t/loaded-torchscript-models-have-wrong-output-shape/6357
**Category:** RLlib
**Created:** [June 1, 2022, 6:21am UTC](https://discuss.ray.io/t/loaded-torchscript-models-have-wrong-output-shape/6357 "2022-06-01T06:21:09Z")
**Posts on this page:** 6
**Page:** 1
### Author:  [@GallantWood](https://discuss.ray.io/u/GallantWood)
#### Post date: [June 1, 2022, 6:21am UTC](https://discuss.ray.io/t/loaded-torchscript-models-have-wrong-output-shape/6357/1 "2022-06-01T06:21:10Z")
**How severe does this issue affect your experience of using Ray?**
- High: It blocks me to complete my task.
I am relatively new to RLlib. I have restored a checkpoint and saved my (non RNN) custom DQN model for inference using export\_model. This generates a torchscript model (model.pt). One subtlety is that this model expects tensor inputs for (state, seq\_lens) so ([], None) will throw an error. However, if an appropriate dummy input is provided, we can do a forward pass.
The problem is that the code:
> model = torch.jit.load(model\_path)
> model.eval()
> obs = torch.rand((1, 128, 111))
> obs\_dict = {“obs”: obs}
> state = np.array([1.0])
> state = [torch.from\_numpy(state)]
> seq\_lens = np.ndarray((1,))
> seq\_lens = torch.from\_numpy(seq\_lens)
> model(obs\_dict, state, seq\_lens)
has the wrong output shape, in this case 256, and not the model output shape. Can anyone direct me towards the solution to this?
If it helps, here are the final layers in print(model):
> (to\_logits): RecursiveScriptModule(
> original\_name=Sequential
> (0): RecursiveScriptModule(original\_name=LayerNorm)
> (1): RecursiveScriptModule(original\_name=Linear)
> )
> (value\_branch): RecursiveScriptModule(
> original\_name=Sequential
> (0): RecursiveScriptModule(original\_name=LayerNorm)
> (1): RecursiveScriptModule(original\_name=Linear)
> )
> (advantage\_module): RecursiveScriptModule(
> original\_name=Sequential
> (dueling\_A\_0): RecursiveScriptModule(
> original\_name=SlimFC
> (\_model): RecursiveScriptModule(
> original\_name=Sequential
> (0): RecursiveScriptModule(original\_name=Linear)
> (1): RecursiveScriptModule(original\_name=ReLU)
> )
> )
> (A): RecursiveScriptModule(
> original\_name=SlimFC
> (\_model): RecursiveScriptModule(
> original\_name=Sequential
> (0): RecursiveScriptModule(original\_name=Linear)
> )
> )
> )
and this is the output to print(model.code):
> def forward(self,
> input\_dict: Dict[str, Tensor],
> state: List[Tensor],
> seq\_lens: Tensor) → Tuple[Tensor, List[Tensor]]:
> to\_logits = self.to\_logits
> hidden\_layers = self.hidden\_layers
> obs = input\_dict[“obs”]
> \_0, = state
> \_1 = (to\_logits).forward((hidden\_layers).forward(obs, ), )
> return (\_1, [\_0])
Since this whole method is used for inference, shouldn’t the output of
> model(obs\_dict, state, seq\_lens)
be of shape num\_outputs where I guess argmax corresponds exactly to policy.compute\_single\_action(obs)[0] ?
In other words, how can I get this single action output through a forward pass of a saved torchscript model?
---
### Author:  [@rliaw](https://discuss.ray.io/u/rliaw)
#### Post date: [June 2, 2022, 9:47pm UTC](https://discuss.ray.io/t/loaded-torchscript-models-have-wrong-output-shape/6357/2 "2022-06-02T21:47:32Z")
Uh oh, that’s not great. @gjoliver is working on improving the inference path. Can someone on RLlib side take a look?
BTW @GallantWood do you think you could provide a copy of the actual model somewhere? i.e., load the model\_path somewhere we can download?
---
### Author:  [@gjoliver](https://discuss.ray.io/u/gjoliver)
#### Post date: [June 2, 2022, 10:36pm UTC](https://discuss.ray.io/t/loaded-torchscript-models-have-wrong-output-shape/6357/3 "2022-06-02T22:36:02Z")
I suspect what you get are the raw logits.
RLlib policies have build-in action distribution functions to sample an action from the distributions constructed based on these raw logit outputs.
can you open a new github issue with a stripped down script demonstrating this problem?
with the actual model config, we can help make sure.
btw, we will be providing utils to make it much easier to run a checkpointed policy. having to come up with dummy state and seq\_lens is painful.
---
### Author:  [@GallantWood](https://discuss.ray.io/u/GallantWood)
#### Post date: [June 3, 2022, 7:43am UTC](https://discuss.ray.io/t/loaded-torchscript-models-have-wrong-output-shape/6357/4 "2022-06-03T07:43:38Z")
Thank you @rliaw @gjoliver
I have built a simple reproducible example using cartpole. I will open a github issue with the code.
---
### Author:  [@GallantWood](https://discuss.ray.io/u/GallantWood)
#### Post date: [June 7, 2022, 7:40am UTC](https://discuss.ray.io/t/loaded-torchscript-models-have-wrong-output-shape/6357/5 "2022-06-07T07:40:02Z")
@gjoliver @rliaw
Here is a link to the issue with complete script demonstrating the problem:
>
>
> \### What happened + What you expected to happen
>
> I have restored a checkpoint an…d saved my (non RNN) custom DQN model for inference using export\_model. This generates a torchscript model (model.pt). One subtlety is that this model expects tensor inputs for (state, seq\_lens) so (\[\], None) will throw an error. However, if an appropriate dummy input is provided, we can do a forward pass.
>
> The issue is that this forward pass returns the wrong output shape, in this case 256, and not the model output shape.
>
> \### Versions / Dependencies
>
> Ray 1.12.1
>
> \### Reproduction script
>
> Here is a reproducible example for Cartpole.
> 1. Train: train a custom DQN model.
> 2. Save: restore a checkpoint and export the Torchscript model.
> 3. Test: use the saved Torchscript model for inference.
>
> Train:
>
> \`\`\`python
> import torch
> from torch import nn
> import ray
> from ray import tune
> from ray.rllib.models import ModelCatalog
> from ray.rllib.models.torch.torch\_modelv2 import TorchModelV2
>
> class MyCustomModel(TorchModelV2, nn.Module):
>
> def \_\_init\_\_(self, obs\_space, action\_space, num\_outputs, model\_config, name, \*\*kwargs):
>
> TorchModelV2.\_\_init\_\_(self, obs\_space, action\_space, num\_outputs, model\_config, name)
> nn.Module.\_\_init\_\_(self)
>
> self.hidden\_layers = nn.Sequential(
> nn.Linear(4, 256),
> nn.Linear(256, num\_outputs)
> )
> self.to\_logits = nn.Linear(256, num\_outputs)
> self.value\_branch = nn.Linear(256, 1)
> self.\_output = None
>
> def forward(self, input\_dict, state, seq\_lens):
> inputs = input\_dict\["obs"\]
> self.\_output = self.hidden\_layers(inputs)
> logits = self.to\_logits(self.\_output)
> return logits, state
>
> def value\_function(self):
> value\_out = self.value\_branch(self.\_output)
> return torch.reshape(value\_out, \[-1\])
>
>
> if \_\_name\_\_ == "\_\_main\_\_":
>
> ray.shutdown()
> ray.init()
>
> ModelCatalog.register\_custom\_model('my\_custom\_model', MyCustomModel)
>
> config = {
> "env": "CartPole-v0",
> "model": {
> "custom\_model": "my\_custom\_model",
> },
> "framework": "torch",
> "num\_gpus": 0,
> "num\_workers": 1,
> }
>
> stop = {"timesteps\_total": 1e4}
>
> analysis = tune.run(
> "DQN",
> config = config,
> stop = stop,
> checkpoint\_at\_end = True
> )
> \`\`\`
>
> Save:
>
> \`\`\`python
> import ray
> from ray import tune
> from ray.rllib.agents import dqn
> from ray.rllib.models import ModelCatalog
>
> from cartpole\_dqn\_train import MyCustomModel
>
> if \_\_name\_\_ == '\_\_main\_\_':
>
> ray.shutdown()
> ray.init()
>
> checkpoint\_path = "DQN/DQN\_CartPole-v0\_85778\_00000\_0\_2022-06-03\_09-24-24/checkpoint\_000010/checkpoint-10"
> export\_path = "saved\_models"
>
> ModelCatalog.register\_custom\_model('my\_custom\_model', MyCustomModel)
>
> config = {
> "env": "CartPole-v0",
> "model": {
> "custom\_model": "my\_custom\_model",
> },
> "framework": "torch",
> "num\_gpus": 0,
> "num\_workers": 1,
> }
>
> agent = dqn.DQNTrainer(config = config)
> agent.restore(checkpoint\_path)
> policy = agent.get\_policy()
> policy.export\_model(export\_path)
> \`\`\`
>
> Test:
>
> \`\`\`python
> import os
> import numpy as np
> import torch
>
> if \_\_name\_\_ == '\_\_main\_\_':
>
> model\_path = 'saved\_models/model.pt'
>
> model = torch.jit.load(model\_path)
> model.eval()
>
> obs = torch.rand((1, 4))
> obs\_dict = {"obs": obs}
>
> state = np.array(\[1.0\])
> state = \[torch.from\_numpy(state)\]
>
> seq\_lens = np.ndarray((1,))
> seq\_lens = torch.from\_numpy(seq\_lens)
>
> print(model(obs\_dict, state, seq\_lens)\[0\].shape) # = torch.Size(\[1, 256\])
> \`\`\`
>
> \### Issue Severity
>
> High: It blocks me from completing my task.
---
### Author:  [@gjoliver](https://discuss.ray.io/u/gjoliver)
#### Post date: [June 7, 2022, 4:15pm UTC](https://discuss.ray.io/t/loaded-torchscript-models-have-wrong-output-shape/6357/6 "2022-06-07T16:15:32Z")
ok cool, let’s move our discussion there.