Inductor stable baselines assertion errors
Labels: triaged, oncall: pt2

Followup on pytorch/torchdynamo#797 

`pip install stable-baselines3[extra]`

## Repro

```python
from stable_baselines3 import PPO
import torchdynamo 



@torchdynamo.optimize("inductor")
def train():
    model = PPO("MlpPolicy", "CartPole-v1").learn(10_000)


import time 
tic = time.time()
train()
toc = time.time()
print(toc - tic)
```

Lots of assertion errors



## Logs

```
(dynamo) ubuntu@ip-172-31-31-152:~/tests$ python baseline.py 
torchdynamo.symbolic_convert: [WARNING] Graph break: call_method GetAttrVariable(UserDefinedObjectVariable(PPO), policy_aliases) __contains__ [ConstantVariable(str)] {} from user code at   File "baseline.py", line 8, in train
    model = PPO("MlpPolicy", "CartPole-v1").learn(10_000)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/ppo/ppo.py", line 102, in __init__
    super().__init__(
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 75, in __init__
    super().__init__(
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/base_class.py", line 106, in __init__
    self.policy_class = self._get_policy_from_name(policy)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/base_class.py", line 340, in _get_policy_from_name
    if policy_name in self.policy_aliases:

torchdynamo.convert_frame: [ERROR] WON'T CONVERT get_device /opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/utils.py line 134 
due to: 
Traceback (most recent call last):
  File "/home/ubuntu/torchdynamo/torchdynamo/variables/tensor.py", line 268, in create
    assert (
AssertionError: torch.* op returned non-Tensor bool call_function <function is_available at 0x7f3904cc4310>

from user code:
   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/utils.py", line 151, in get_device
    if device.type == th.device("cuda").type and not th.cuda.is_available():

Set torchdynamo.config.verbose=True for more information
==========
torchdynamo.symbolic_convert: [WARNING] Graph break: numpy from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/gym/envs/classic_control/cartpole.py", line 83, in __init__
    np.finfo(np.float32).max,

torchdynamo.symbolic_convert: [WARNING] Graph break: numpy from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/gym/envs/classic_control/cartpole.py", line 85, in <graph break in __init__>
    np.finfo(np.float32).max,

torchdynamo.symbolic_convert: [WARNING] Graph break: numpy from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/gym/envs/classic_control/cartpole.py", line 80, in <graph break in __init__>
    high = np.array(

torchdynamo.symbolic_convert: [WARNING] Graph break: numpy from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/gym/envs/classic_control/cartpole.py", line 90, in <graph break in __init__>
    self.action_space = spaces.Discrete(2)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/gym/spaces/discrete.py", line 17, in __init__
    super(Discrete, self).__init__((), np.int64, seed)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/gym/spaces/space.py", line 23, in __init__
    self.dtype = None if dtype is None else np.dtype(dtype)

torchdynamo.convert_frame: [ERROR] WON'T CONVERT <graph break in __init__> /opt/conda/envs/dynamo/lib/python3.8/site-packages/gym/envs/classic_control/cartpole.py line 90 
due to: 
Traceback (most recent call last):
  File "/home/ubuntu/torchdynamo/torchdynamo/variables/functions.py", line 34, in wrap_bound_arg
    assert isinstance(val, VariableTracker), typestr(val)
AssertionError: type

from user code:
   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/gym/envs/classic_control/cartpole.py", line 91, in <graph break in __init__>
    self.observation_space = spaces.Box(-high, high, dtype=np.float32)

Set torchdynamo.config.verbose=True for more information
==========
torchdynamo.convert_frame: [ERROR] WON'T CONVERT __init__ /opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/policies.py line 413 
due to: 
Traceback (most recent call last):
  File "/home/ubuntu/torchdynamo/torchdynamo/variables/functions.py", line 34, in wrap_bound_arg
    assert isinstance(val, VariableTracker), typestr(val)
AssertionError: type

from user code:
   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/policies.py", line 440, in __init__
    super().__init__(

Set torchdynamo.config.verbose=True for more information
==========
torchdynamo.convert_frame: [ERROR] WON'T CONVERT __init__ /opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/policies.py line 269 
due to: 
Traceback (most recent call last):
  File "/home/ubuntu/torchdynamo/torchdynamo/variables/functions.py", line 34, in wrap_bound_arg
    assert isinstance(val, VariableTracker), typestr(val)
AssertionError: type

from user code:
   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/policies.py", line 270, in __init__
    super().__init__(*args, **kwargs)

Set torchdynamo.config.verbose=True for more information
==========
torchdynamo.symbolic_convert: [WARNING] Graph break: Patched init cannot be inlined. from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/policies.py", line 70, in __init__
    super().__init__()

torchdynamo.symbolic_convert: [WARNING] Graph break: call_function in skip_files /opt/conda/envs/dynamo/lib/python3.8/functools.py from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/torch_layers.py", line 44, in __init__
    super().__init__(observation_space, get_flattened_obs_dim(observation_space))
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/preprocessing.py", line 177, in get_flattened_obs_dim
    return spaces.utils.flatdim(observation_space)

torchdynamo.symbolic_convert: [WARNING] Graph break: Patched init cannot be inlined. from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/torch_layers.py", line 22, in __init__
    super().__init__()

torchdynamo.symbolic_convert: [WARNING] Graph break: Patched init cannot be inlined. from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/policies.py", line 542, in _build
    self._build_mlp_extractor()
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/policies.py", line 528, in _build_mlp_extractor
    self.mlp_extractor = MlpExtractor(
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/torch_layers.py", line 172, in __init__
    super().__init__()

torchdynamo.symbolic_convert: [WARNING] Graph break: Patched init cannot be inlined. from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/torch/nn/modules/container.py", line 85, in __init__
    super(Sequential, self).__init__()

torchdynamo.symbolic_convert: [WARNING] Graph break: construct nn.Module: Linear from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/policies.py", line 555, in <graph break in _build>
    self.action_net = self.action_dist.proba_distribution_net(latent_dim=latent_dim_pi)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/distributions.py", line 271, in proba_distribution_net
    action_logits = nn.Linear(latent_dim, self.action_dim)

torchdynamo.symbolic_convert: [WARNING] Graph break: construct nn.Module: Linear from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/policies.py", line 559, in <graph break in _build>
    self.value_net = nn.Linear(self.mlp_extractor.latent_dim_vf, 1)

torchdynamo.symbolic_convert: [WARNING] Graph break: numpy from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/policies.py", line 568, in <graph break in _build>
    self.features_extractor: np.sqrt(2),

torchdynamo.symbolic_convert: [WARNING] Graph break: numpy from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/policies.py", line 569, in <graph break in _build>
    self.mlp_extractor: np.sqrt(2),

torchdynamo.convert_frame: [ERROR] WON'T CONVERT <graph break in _build> /opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/policies.py line 569 
due to: 
Traceback (most recent call last):
  File "/home/ubuntu/torchdynamo/torchdynamo/symbolic_convert.py", line 891, in BUILD_MAP
    assert isinstance(k, ConstantVariable) or (
AssertionError

from user code:
   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/policies.py", line 567, in <graph break in _build>
    module_gains = {

Set torchdynamo.config.verbose=True for more information
==========
torchdynamo.symbolic_convert: [WARNING] Graph break: hasattr: GetAttrVariable(UserDefinedClassVariable(), step) from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/torch/optim/optimizer.py", line 49, in __init__
    self._hook_for_profile()
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/torch/optim/optimizer.py", line 130, in _hook_for_profile
    hooked = getattr(self.__class__.step, "hooked", None)

torchdynamo.symbolic_convert: [WARNING] Graph break: call_function in skip_files /opt/conda/envs/dynamo/lib/python3.8/collections/__init__.py from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/torch/optim/optimizer.py", line 56, in <graph break in __init__>
    self.state = defaultdict(dict)

torchdynamo.symbolic_convert: [WARNING] Graph break: call_function BuiltinVariable(list) [UserDefinedObjectVariable(generator)] {} from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/torch/optim/optimizer.py", line 59, in <graph break in __init__>
    param_groups = list(params)

torchdynamo.symbolic_convert: [WARNING] Graph break: generic_jump GetAttrVariable(TensorVariable(), is_leaf) from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/torch/optim/optimizer.py", line 66, in <graph break in __init__>
    self.add_param_group(param_group)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/torch/optim/optimizer.py", line 309, in add_param_group
    if not self.defaults.get('differentiable', None) and not (param.is_leaf or param.retains_grad):

torchdynamo.symbolic_convert: [WARNING] Graph break: call_function UserDefinedObjectVariable(time) [] {} from user code at   File "baseline.py", line 8, in <graph break in train>
    model = PPO("MlpPolicy", "CartPole-v1").learn(10_000)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/ppo/ppo.py", line 310, in learn
    return super().learn(
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 239, in learn
    total_timesteps, callback = self._setup_learn(
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/base_class.py", line 425, in _setup_learn
    self.start_time = time.time()

torchdynamo.symbolic_convert: [WARNING] Graph break: call_function in skip_files /opt/conda/envs/dynamo/lib/python3.8/collections/__init__.py from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/base_class.py", line 429, in <graph break in _setup_learn>
    self.ep_info_buffer = deque(maxlen=100)

torchdynamo.symbolic_convert: [WARNING] Graph break: call_function in skip_files /opt/conda/envs/dynamo/lib/python3.8/collections/__init__.py from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/base_class.py", line 430, in <graph break in _setup_learn>
    self.ep_success_buffer = deque(maxlen=100)

torchdynamo.symbolic_convert: [WARNING] Graph break: call_method UserDefinedObjectVariable(RandomState) uniform [] {'low': ConstantVariable(float), 'high': ConstantVariable(float), 'size': ConstantVariable(tuple)} from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/base_class.py", line 446, in <graph break in _setup_learn>
    self._last_obs = self.env.reset()  # pytype: disable=annotation-type-mismatch
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/vec_env/dummy_vec_env.py", line 63, in reset
    obs = self.envs[env_idx].reset()
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/monitor.py", line 79, in reset
    return self.env.reset(**kwargs)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/gym/wrappers/time_limit.py", line 27, in reset
    return self.env.reset(**kwargs)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/gym/envs/classic_control/cartpole.py", line 162, in reset
    self.state = self.np_random.uniform(low=-0.05, high=0.05, size=(4,))

torchdynamo.symbolic_convert: [WARNING] Graph break: numpy from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/base_class.py", line 447, in <graph break in _setup_learn>
    self._last_episode_starts = np.ones((self.env.num_envs,), dtype=bool)

torchdynamo.symbolic_convert: [WARNING] Graph break: call_function in skip_files /opt/conda/envs/dynamo/lib/python3.8/os.py from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/base_class.py", line 459, in <graph break in _setup_learn>
    self._logger = utils.configure_logger(self.verbose, self.tensorboard_log, tb_log_name, reset_num_timesteps)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/utils.py", line 210, in configure_logger
    return configure(save_path, format_strings=format_strings)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/logger.py", line 596, in configure
    folder = os.getenv("SB3_LOGDIR")

torchdynamo.symbolic_convert: [WARNING] Graph break: call_function BuiltinVariable(locals) [] {} from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 243, in <graph break in learn>
    callback.on_training_start(locals(), globals())

torchdynamo.symbolic_convert: [WARNING] Graph break: UnspecializedNNModuleVariable missing train from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 247, in <graph break in learn>
    continue_training = self.collect_rollouts(self.env, callback, self.rollout_buffer, n_rollout_steps=self.n_steps)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 148, in collect_rollouts
    self.policy.set_training_mode(False)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/policies.py", line 215, in set_training_mode
    self.train(mode)

torchdynamo.symbolic_convert: [WARNING] Graph break: numpy from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 151, in <graph break in collect_rollouts>
    rollout_buffer.reset()
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/buffers.py", line 360, in reset
    self.observations = np.zeros((self.buffer_size, self.n_envs) + self.obs_shape, dtype=np.float32)

torchdynamo.symbolic_convert: [WARNING] Graph break: call_function args: NumpyVariable()  from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 165, in <graph break in collect_rollouts>
    obs_tensor = obs_as_tensor(self._last_obs, self.device)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/utils.py", line 451, in obs_as_tensor
    return th.as_tensor(obs).to(device)

torchinductor.ir: [WARNING] DeviceCopy
torchinductor.compile_fx: [WARNING] skipping cudagraphs due to multiple devices
torchdynamo.symbolic_convert: [WARNING] /opt/conda/envs/dynamo/lib/python3.8/site-packages/torch/distributions/categorical.py <function Categorical.sample at 0x7f3903e27b80> [] {} missing a required argument: 'self'
torchdynamo.symbolic_convert: [WARNING] Graph break: arg mismatch inlining from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 166, in <graph break in collect_rollouts>
    actions, values, log_probs = self.policy(obs_tensor)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/policies.py", line 593, in forward
    actions = distribution.get_actions(deterministic=deterministic)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/distributions.py", line 80, in get_actions
    return self.sample()
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/distributions.py", line 285, in sample
    return self.distribution.sample()

torchdynamo.symbolic_convert: [WARNING] /opt/conda/envs/dynamo/lib/python3.8/site-packages/torch/distributions/categorical.py <function Categorical.sample at 0x7f3903e27b80> [] {} missing a required argument: 'self'
torchdynamo.symbolic_convert: [WARNING] Graph break: call_function BuiltinVariable(islice) [UserDefinedObjectVariable(odict_values), ConstantVariable(int), ConstantVariable(NoneType)] {} from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/torch/nn/modules/container.py", line 100, in _get_item_by_idx
    return next(islice(iterator, idx, None))

torchdynamo.symbolic_convert: [WARNING] /opt/conda/envs/dynamo/lib/python3.8/site-packages/torch/distributions/categorical.py <function Categorical.sample at 0x7f3903e27b80> [] {} missing a required argument: 'self'
torchdynamo.symbolic_convert: [WARNING] /opt/conda/envs/dynamo/lib/python3.8/site-packages/torch/distributions/categorical.py <function Categorical.sample at 0x7f3903e27b80> [] {} missing a required argument: 'self'
torchdynamo.symbolic_convert: [WARNING] /opt/conda/envs/dynamo/lib/python3.8/site-packages/torch/distributions/categorical.py <function Categorical.log_prob at 0x7f3903e27c10> [TensorVariable()] {} missing a required argument: 'value'
torchdynamo.symbolic_convert: [WARNING] Graph break: arg mismatch inlining from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/policies.py", line 594, in <graph break in forward>
    log_prob = distribution.log_prob(actions)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/distributions.py", line 279, in log_prob
    return self.distribution.log_prob(actions)

torchdynamo.symbolic_convert: [WARNING] /opt/conda/envs/dynamo/lib/python3.8/site-packages/torch/distributions/categorical.py <function Categorical.log_prob at 0x7f3903e27c10> [TensorVariable()] {} missing a required argument: 'value'
torchdynamo.symbolic_convert: [WARNING] Graph break: Tensor.numpy from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 167, in <graph break in collect_rollouts>
    actions = actions.cpu().numpy()

torchinductor.ir: [WARNING] DeviceCopy
torchinductor.compile_fx: [WARNING] skipping cudagraphs due to multiple devices
torchdynamo.symbolic_convert: [WARNING] Graph break: call_method ConstantVariable(str) __contains__ [GetAttrVariable(GetAttrVariable(ConstantVariable(int64), dtype), char)] {} from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 175, in <graph break in collect_rollouts>
    new_obs, rewards, dones, infos = env.step(clipped_actions)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/vec_env/base_vec_env.py", line 162, in step
    return self.step_wait()
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/vec_env/dummy_vec_env.py", line 43, in step_wait
    obs, self.buf_rews[env_idx], self.buf_dones[env_idx], self.buf_infos[env_idx] = self.envs[env_idx].step(
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/monitor.py", line 90, in step
    observation, reward, done, info = self.env.step(action)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/gym/wrappers/time_limit.py", line 18, in step
    observation, reward, done, info = self.env.step(action)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/gym/envs/classic_control/cartpole.py", line 105, in step
    assert self.action_space.contains(action), err_msg
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/gym/spaces/discrete.py", line 26, in contains
    x.dtype.char in np.typecodes["AllInteger"] and x.shape == ()

torchdynamo.symbolic_convert: [WARNING] Graph break: call_function BuiltinVariable(locals) [] {} from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 180, in <graph break in collect_rollouts>
    callback.update_locals(locals())

torchdynamo.symbolic_convert: [WARNING] Graph break: numpy from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 184, in <graph break in collect_rollouts>
    self._update_info_buffer(infos)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/base_class.py", line 475, in _update_info_buffer
    dones = np.array([False] * len(infos))

torchdynamo.symbolic_convert: [WARNING] Graph break: call_method ConstDictVariable() get [ConstantVariable(str)] {} from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/base_class.py", line 477, in <graph break in _update_info_buffer>
    maybe_ep_info = info.get("episode")

torchdynamo.symbolic_convert: [WARNING] Graph break: call_method ConstDictVariable() get [ConstantVariable(str)] {} from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/base_class.py", line 478, in <graph break in _update_info_buffer>
    maybe_is_success = info.get("is_success")

torchdynamo.symbolic_convert: [WARNING] Graph break: numpy from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 189, in <graph break in collect_rollouts>
    actions = actions.reshape(-1, 1)

torchdynamo.symbolic_convert: [WARNING] Graph break: numpy from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/buffers.py", line 436, in add
    self.observations[self.pos] = np.array(obs).copy()

torchdynamo.symbolic_convert: [WARNING] Graph break: numpy from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/buffers.py", line 437, in <graph break in add>
    self.actions[self.pos] = np.array(action).copy()

torchdynamo.symbolic_convert: [WARNING] Graph break: numpy from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/buffers.py", line 438, in <graph break in add>
    self.rewards[self.pos] = np.array(reward).copy()

torchdynamo.symbolic_convert: [WARNING] Graph break: numpy from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/buffers.py", line 439, in <graph break in add>
    self.episode_starts[self.pos] = np.array(episode_start).copy()

torchdynamo.symbolic_convert: [WARNING] Graph break: Tensor.numpy from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/buffers.py", line 440, in <graph break in add>
    self.values[self.pos] = value.clone().cpu().numpy().flatten()

torchinductor.ir: [WARNING] DeviceCopy
torchinductor.compile_fx: [WARNING] skipping cudagraphs due to multiple devices
Traceback (most recent call last):
  File "baseline.py", line 13, in <module>
    train()
  File "/home/ubuntu/torchdynamo/torchdynamo/eval_frame.py", line 166, in _fn
    return fn(*args, **kwargs)
  File "baseline.py", line 6, in train
    @torchdynamo.optimize("inductor")
  File "baseline.py", line 8, in <graph break in train>
    model = PPO("MlpPolicy", "CartPole-v1").learn(10_000)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/ppo/ppo.py", line 297, in learn
    def learn(
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 239, in learn
    total_timesteps, callback = self._setup_learn(
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 239, in <graph break in learn>
    total_timesteps, callback = self._setup_learn(
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 243, in <graph break in learn>
    callback.on_training_start(locals(), globals())
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 243, in <graph break in learn>
    callback.on_training_start(locals(), globals())
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 243, in <graph break in learn>
    callback.on_training_start(locals(), globals())
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 126, in collect_rollouts
    def collect_rollouts(
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 148, in <graph break in collect_rollouts>
    self.policy.set_training_mode(False)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 151, in <graph break in collect_rollouts>
    rollout_buffer.reset()
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 165, in <graph break in collect_rollouts>
    obs_tensor = obs_as_tensor(self._last_obs, self.device)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 166, in <graph break in collect_rollouts>
    actions, values, log_probs = self.policy(obs_tensor)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 167, in <graph break in collect_rollouts>
    actions = actions.cpu().numpy()
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 175, in <graph break in collect_rollouts>
    new_obs, rewards, dones, infos = env.step(clipped_actions)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 180, in <graph break in collect_rollouts>
    callback.update_locals(locals())
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 180, in <graph break in collect_rollouts>
    callback.update_locals(locals())
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 184, in <graph break in collect_rollouts>
    self._update_info_buffer(infos)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 189, in <graph break in collect_rollouts>
    actions = actions.reshape(-1, 1)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/on_policy_algorithm.py", line 204, in <graph break in collect_rollouts>
    rollout_buffer.add(self._last_obs, actions, rewards, self._last_episode_starts, values, log_probs)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/buffers.py", line 408, in add
    def add(
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/buffers.py", line 436, in <graph break in add>
    self.observations[self.pos] = np.array(obs).copy()
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/buffers.py", line 436, in <graph break in add>
    self.observations[self.pos] = np.array(obs).copy()
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/buffers.py", line 436, in <graph break in add>
    self.observations[self.pos] = np.array(obs).copy()
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/buffers.py", line 437, in <graph break in add>
    self.actions[self.pos] = np.array(action).copy()
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/buffers.py", line 437, in <graph break in add>
    self.actions[self.pos] = np.array(action).copy()
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/buffers.py", line 437, in <graph break in add>
    self.actions[self.pos] = np.array(action).copy()
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/buffers.py", line 438, in <graph break in add>
    self.rewards[self.pos] = np.array(reward).copy()
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/buffers.py", line 438, in <graph break in add>
    self.rewards[self.pos] = np.array(reward).copy()
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/buffers.py", line 438, in <graph break in add>
    self.rewards[self.pos] = np.array(reward).copy()
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/buffers.py", line 439, in <graph break in add>
    self.episode_starts[self.pos] = np.array(episode_start).copy()
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/buffers.py", line 439, in <graph break in add>
    self.episode_starts[self.pos] = np.array(episode_start).copy()
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/stable_baselines3/common/buffers.py", line 439, in <graph break in add>
    self.episode_starts[self.pos] = np.array(episode_start).copy()
RuntimeError: Can't call numpy() on Tensor that requires grad. Use tensor.detach().numpy() instead.
```

cc @ezyang @soumith @wconstab @ngimel @bdhirsh