Method tensor_to_arg

examples/diffusion/python_flux/models.py:126–130

(self, tensor)

        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],
                              lens=list(tensor.size()),
                              strides=list(tensor.stride()))
        return mgx.argument_from_pointer(mgx_shape, tensor.data_ptr())

    def prealloc_buffers(self, param_names):
        for param_name in param_names: