shingjan commented on code in PR #13623:
URL: https://github.com/apache/tvm/pull/13623#discussion_r1050009039
##########
python/tvm/contrib/graph_executor.py:
##########
@@ -195,6 +198,49 @@ def set_input(self, key=None, value=None, **params):
if val:
self._get_input(k).copyfrom(params[k])
+ def set_output(self, key, value):
+ """Set outputs to the module
+
+ Parameters
+ ----------
+ key : int or str
+ The output key
+
+ value : the output value
+ The output value
+ """
+ self._set_output(key, value)
+
+ def set_input_zero_copy(self, key, value, **params):
+ """Set inputs to the module via kwargs with zero memory copy
+
+ Parameters
+ ----------
+ key : int or str
+ The input key
+
+ value : the input value in DLPack
+ The input key
+
+ params : dict of str to NDArray
+ Additional arguments
+ """
+ self._set_input_zero_copy(key, value)
+ self.set_input(None, None, **params)
Review Comment:
good point. Updated this method to use zero copy.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
