# IMPALA agent not working

**URL:** <https://discuss.ray.io/t/impala-agent-not-working/8507>\
**Category:** RLlib\
**Created:** [December 1, 2022, 3:10pm UTC](https://discuss.ray.io/t/impala-agent-not-working/8507 "2022-12-01T15:10:25Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![aea](https://avatars.discourse-cdn.com/v4/letter/a/e56c9b/32.png) [@aea](https://discuss.ray.io/u/aea)\
**Post date:** [December 1, 2022, 3:10pm UTC](https://discuss.ray.io/t/impala-agent-not-working/8507/1 "2022-12-01T15:10:25Z")

</div>

Hey everybody, so I’m using the **rayproject/ray-ml:b42155-py310-cu116** docker image and I’m trying to run the IMPALA agent as described in the documentation [RLlib IMPALA documentation](https://docs.ray.io/en/master/rllib/rllib-algorithms.html#importance-weighted-actor-learner-architecture-impala)

```auto
from ray.rllib.algorithms.impala import ImpalaConfig
config = ImpalaConfig()
config = config.training(lr=0.0003, train_batch_size=512) # doctest: +SKIP
config = config.resources(num_gpus=4) # doctest: +SKIP
config = config.rollouts(num_rollout_workers=64) # doctest: +SKIP
print(config.to_dict()) # doctest: +SKIP
# Build a Algorithm object from the config and run 1 training iteration.
algo = config.build(env="CartPole-v1") # doctest: +SKIP
algo.train() # doctest: +SKIP

```

When I run this with 1 gpu (changing num\_gpus=1) It crashes  
`INVALID_ARGUMENT: Input to reshape is a tensor with 512 values, but the requested shae has 500`  
Does anyone have any idea why is this not working out of the box ?

---

<div class="post-metadata">

**Author:** ![kourosh](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/kourosh/32/2130_2.png) [@kourosh](https://discuss.ray.io/u/kourosh)\
**Post date:** [January 9, 2023, 2:36am UTC](https://discuss.ray.io/t/impala-agent-not-working/8507/2 "2023-01-09T02:36:05Z")

</div>

Hi @aea,

The docstring example codes are not guaranteed to work right of the box as of now since they are not doc-tested. If you don’t provide `train_batch_size` to leave the default 500, your code would work. Basically:

```auto
from ray.rllib.algorithms.impala import ImpalaConfig
config = ImpalaConfig()
config = config.training(lr=0.0003) # doctest: +SKIP
config = config.resources(num_gpus=1) # doctest: +SKIP
config = config.rollouts(num_rollout_workers=64) # doctest: +SKIP
print(config.to_dict()) # doctest: +SKIP
# Build a Algorithm object from the config and run 1 training iteration.
algo = config.build(env="CartPole-v1") # doctest: +SKIP
algo.train() # doctest: +SKIP

```