Method load

deepspeed/compile/util.py:215–216

(self, name: str)

        self.env_values[name] = map_aggregate(v, lambda x: self._to_cpu(x) if offload else x)

    def load(self, name: str) -> Argument:
        return map_aggregate(self.env_values[name], lambda x: self._from_cpu(x))

    def get_offloaded_value(self, name: str) -> Argument:
        return self.env_values[name]