# Buggy code
def action_prob(self, batch: SampleBatchType) -> TensorType:
    """Returns the probs for the batch actions for the current policy."""

    num_state_inputs = 0
    for k in batch.keys():
        if k.startswith("state_in_"):
            num_state_inputs += 1
    state_keys = ["state_in_{}".format(i) for i in range(num_state_inputs)]
    log_likelihoods = self.policy.compute_log_likelihoods(
        actions=batch[SampleBatch.ACTIONS],
        obs_batch=batch[SampleBatch.CUR_OBS],
        state_batches=[batch[k] for k in state_keys],
        prev_action_batch=batch.data.get(SampleBatch.PREV_ACTIONS),
        prev_reward_batch=batch.data.get(SampleBatch.PREV_REWARDS),
    )
    return log_likelihoods


# Traceback (AttributeError)
Traceback (most recent call last):
File "C:\Users\Julius\Anaconda3\envs\ray\lib\site-packages\ray\tune\trial_runner.py", line 497, in _process_trial
result = self.trial_executor.fetch_result(trial)
File "C:\Users\Julius\Anaconda3\envs\ray\lib\site-packages\ray\tune\ray_trial_executor.py", line 434, in fetch_result
result = ray.get(trial_future[0], DEFAULT_GET_TIMEOUT)
File "C:\Users\Julius\Anaconda3\envs\ray\lib\site-packages\ray\worker.py", line 1553, in get
raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(AttributeError): ray::MARWIL.train() (pid=9136, ip=10.0.0.18)
File "python\ray\_raylet.pyx", line 474, in ray._raylet.execute_task
File "python\ray\_raylet.pyx", line 427, in ray._raylet.execute_task.function_executor
File "C:\Users\Julius\Anaconda3\envs\ray\lib\site-packages\ray\function_manager.py", line 567, in actor_method_executor
raise e
File "C:\Users\Julius\Anaconda3\envs\ray\lib\site-packages\ray\function_manager.py", line 559, in actor_method_executor
method_returns = method(actor, *args, **kwargs)
File "C:\Users\Julius\Anaconda3\envs\ray\lib\site-packages\ray\rllib\agents\trainer.py", line 522, in train
raise e
File "C:\Users\Julius\Anaconda3\envs\ray\lib\site-packages\ray\rllib\agents\trainer.py", line 508, in train
result = Trainable.train(self)
File "C:\Users\Julius\Anaconda3\envs\ray\lib\site-packages\ray\tune\trainable.py", line 337, in train
result = self.step()
File "C:\Users\Julius\Anaconda3\envs\ray\lib\site-packages\ray\rllib\agents\trainer_template.py", line 110, in step
res = next(self.train_exec_impl)
File "C:\Users\Julius\Anaconda3\envs\ray\lib\site-packages\ray\util\iter.py", line 758, in __next__
return next(self.built_iterator)
File "C:\Users\Julius\Anaconda3\envs\ray\lib\site-packages\ray\util\iter.py", line 793, in apply_foreach
result = fn(item)
File "C:\Users\Julius\Anaconda3\envs\ray\lib\site-packages\ray\rllib\execution\metric_ops.py", line 87, in __call__
res = summarize_episodes(episodes, orig_episodes)
File "C:\Users\Julius\Anaconda3\envs\ray\lib\site-packages\ray\rllib\evaluation\metrics.py", line 173, in summarize_episodes
metrics[k] = np.mean(v_list)
File "<__array_function__ internals>", line 6, in mean
File "C:\Users\Julius\Anaconda3\envs\ray\lib\site-packages\numpy\core\fromnumeric.py", line 3335, in mean
out=out, **kwargs)
File "C:\Users\Julius\Anaconda3\envs\ray\lib\site-packages\numpy\core\_methods.py", line 161, in _mean
ret = ret.dtype.type(ret / rcount)
AttributeError: 'torch.dtype' object has no attribute 'type'

# Fixed code
def action_prob(self, batch: SampleBatchType) -> np.ndarray:
    """Returns the probs for the batch actions for the current policy."""

    num_state_inputs = 0
    for k in batch.keys():
        if k.startswith("state_in_"):
            num_state_inputs += 1
    state_keys = ["state_in_{}".format(i) for i in range(num_state_inputs)]
    log_likelihoods: TensorType = self.policy.compute_log_likelihoods(
        actions=batch[SampleBatch.ACTIONS],
        obs_batch=batch[SampleBatch.CUR_OBS],
        state_batches=[batch[k] for k in state_keys],
        prev_action_batch=batch.data.get(SampleBatch.PREV_ACTIONS),
        prev_reward_batch=batch.data.get(SampleBatch.PREV_REWARDS),
    )
    return convert_to_numpy(log_likelihoods)