update demo

texts/getting_my_agent_evaluated.md

@@ -93,7 +93,7 @@ class Agent(nn.Module):
 agent = Agent(policy) # instantiate the agent
 
 # A few tests to check if the agent is working
-observations = torch.tensor(env.observation_space.sample()).unsqueeze(0)  # dummy batch of observations
+observations = torch.randn(env.observation_space.shape).unsqueeze(0)  # dummy batch of observations
 actions = agent(observations)
 actions = actions.numpy()[0]
 assert env.action_space.contains(actions)
@@ -109,10 +109,9 @@ from huggingface_hub import metadata_save, HfApi
 
 # Save model along with its card
 metadata_save("model_card.md", {"tags": ["reinforcement-learning", env_id]})
-dummy_input = torch.tensor(env.observation_space.sample()).unsqueeze(0)  # dummy batch of observations
+dummy_input = torch.randn(env.observation_space.shape).unsqueeze(0)  # dummy batch of observations
 agent = torch.jit.trace(agent.eval(), dummy_input)
-agent = torch.jit.freeze(agent)  # required for for the model not to depend on the training library
-agent = torch.jit.optimize_for_inference(agent)
+agent = torch.jit.freeze(agent)  # required for the model not to depend on the training library
 torch.jit.save(agent, "agent.pt")
 
 # Upload model and card to the 🤗 Hub