[GitHub] [incubator-mxnet] samskalicky commented on a change in pull request #19469: Extension bug fixes

[GitBox <gi...@apache.org>]

samskalicky commented on a change in pull request #19469:
URL: https://github.com/apache/incubator-mxnet/pull/19469#discussion_r518444379



##########
File path: python/mxnet/gluon/block.py
##########
@@ -1029,14 +1029,35 @@ def _build_cache(self, *args):
 
         arg_dict, aux_dict = dict(), dict()
         if self._backend:
-            ctx = args[0].context
+            # set context for inputs
+            _, _, ctx_set, _ = _gather_type_ctx_info(list(args))
+            ctx = ctx_set.pop() if len(ctx_set) > 0 else None
             # get list of params in the order of out.list_arguments
-            arg_dict.update({name:args[data_names[name]] if name in data_names.keys() else params[name].data()
-                             for name in out.list_arguments()})
-            aux_dict.update({name:args[data_names[name]] if name in data_names.keys() else params[name].data()
-                             for name in out.list_auxiliary_states()})
-            # Partition the graph.
-            out = out.optimize_for(self._backend, arg_dict, aux_dict, ctx, **self._backend_opts)
+            input_shapes = dict()
+            for name in out.list_arguments():
+                if name in data_names.keys() and data_names[name] < len(args):

Review comment:
       not if the user doesnt specify all the args for a model. if they only specify a few it might be less. basically we're checking to see if all the args for the model were specified by the user. 




----------------------------------------------------------------
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.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org