-
Notifications
You must be signed in to change notification settings - Fork 52
Description
KeyError Traceback (most recent call last)
in
32
33 model_trt = TRTModule()
---> 34 model_trt.load_state_dict(torch.load(OPTIMIZED_MODEL))
/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in load_state_dict(self, state_dict, strict)
1466 load(child, prefix + name + '.')
1467
-> 1468 load(self)
1469 del load
1470
/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in load(module, prefix)
1461 local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
1462 module._load_from_state_dict(
-> 1463 state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs)
1464 for name, child in module._modules.items():
1465 if child is not None:
/usr/local/lib/python3.6/dist-packages/torch2trt-0.4.0-py3.6-linux-aarch64.egg/torch2trt/torch2trt.py in _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
585 error_msgs,
586 ):
--> 587 engine_bytes = state_dict[prefix + "engine"]
588
589 with trt.Logger() as logger, trt.Runtime(logger) as runtime:
KeyError: 'engine'