add online trt export

This commit is contained in:
lyuxiang.lx
2025-01-10 13:55:05 +08:00
parent 426c4001ca
commit 1cfc5dd077
13 changed files with 100 additions and 167 deletions

View File

@@ -134,12 +134,12 @@ class ConditionalCFM(BASECFM):
self.estimator.set_input_shape('cond', (2, 80, x.size(2)))
# run trt engine
self.estimator.execute_v2([x.contiguous().data_ptr(),
mask.contiguous().data_ptr(),
mu.contiguous().data_ptr(),
t.contiguous().data_ptr(),
spks.contiguous().data_ptr(),
cond.contiguous().data_ptr(),
x.data_ptr()])
mask.contiguous().data_ptr(),
mu.contiguous().data_ptr(),
t.contiguous().data_ptr(),
spks.contiguous().data_ptr(),
cond.contiguous().data_ptr(),
x.data_ptr()])
return x
def compute_loss(self, x1, mask, mu, spks=None, cond=None):