def METHOD_NAME(self, x, indices):
    if isinstance(x, (dict, TensorDict)):
        res = TensorDict()
        for key, value in x.items():
            res[key] = self.METHOD_NAME(value, indices)
        return res
    else:
        t = x[indices]
        return t