I am following this tutorial:
https://docs.tvm.ai/tutorials/autotvm/tune_nnvm_cuda.html
I am having trouble implementing the get_network method.
For CPU i am using the following to load from mxnet but for GPU i am confused.
def get_network(name, batch_size):
prefix,epoch = “/opt/models/model”,0
sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)
opt_level = 3
nnvm_sym, nnvm_params = nnvm.frontend.from_mxnet(sym, arg_params, aux_params)
input_shape = (batch_size, 3, 224, 224)
output_shape = (batch_size, 1024)
return nnvm_sym, nnvm_params, input_shape, output_shape
What is similar code for GPU ?