Method _from_cpu

deepspeed/compile/util.py:207–210

(self, v)

        return v

    def _from_cpu(self, v):
        if torch.is_tensor(v) and v in self.original_device:
            return v.to(self.original_device[v])
        return v

    def save(self, name: str, v: Argument, offload) -> None:
        self.env_values[name] = map_aggregate(v, lambda x: self._to_cpu(x) if offload else x)