There doesn’t seem to be anything wrong in the showed part. I think you may try to decrease the the local /global replay buffer size in your implementation, /or skip learning & rolling out by appending radomly generated tensor data into your buffer and see when they will overflow, and track down the distributed bug.
If you would like to have a correct implementation of IMPALA as reference, I have written a tested one before.
Doc here:
https://machin.readthedocs.io/en/latest/tutorials/unleash_distributed_power.html#impala 1
and source code here:

  
      github.com
  
  
    iffiX/machin/blob/master/machin/frame/algorithms/impala.py 3
from typing import Union, Dict, List, Tuple, Callable, Any
import numpy as np
import torch as t
import torch.nn as nn

from machin.frame.buffers.buffer_d import Transition, DistributedBuffer
from machin.model.nets.base import NeuralNetworkModule
from .base import TorchFramework
from .utils import safe_call

from machin.parallel.server import PushPullModelServer
from machin.parallel.distributed import RpcGroup


def _make_tensor_from_batch(batch: List[Any], device, concatenate):
    """
    Used to convert compact every attribute of every step of a whole episode
    into a single tensor.

    Args:


  This file has been truncated. show original