Method save_model

examples/diffusion/python_flux/models.py:122–124

(self, path)

        return {p: self.torch_buffers[p] for p in self.output_names}

    def save_model(self, path):
        os.makedirs(os.path.dirname(path), exist_ok=True)
        mgx.save(self.model, path, format="msgpack")

    def tensor_to_arg(self, tensor):
        mgx_shape = mgx.shape(type=self.torch_to_mgx_dtype_dict[tensor.dtype],