you can try just shoving that into the forward loop of the model, not sure if that would work.
note: dont do something like:
if train_flag:
code
else:
with torch.no_grad
code
because the if statement will break the tracing.
alternatively you can try:
with torch.no_grad()
rknn.load_pytorch(…)