[Question & FeatureRequest] Is that possible to convert a torch.nn.Module to TRTMoudle?
ElinLiu0 opened this issue · comments
Elin commented
For example,i got a module class which based on torch.nn.Moudle()
like this:
model = VALLE(
N_DIM,
NUM_HEAD,
NUM_LAYERS,
norm_first=True,
add_prenet=False,
prefix_mode=PREFIX_MODE,
share_embedding=True,
nar_scale_factor=1.0,
prepend_bos=True,
num_quantizers=NUM_QUANTIZERS,
).to(device)
Is that possible convert this module class into a TRTMoudle individually,and still can be using the methods and attributes from VALLE?
I am seen that TRTMoudle has a function load_state_dict()
,will it have a simple convert like somehow,uh..
TRTModule.from_torch()
or something?
Elin commented
And then just seperate the module initailize and state_dict() loading.