RuntimeError exporting stateless5 model using jit

[Cescfangs]

I'm using ./pruned_transducer_stateless5/export.py with --jit true and get runtime error:

RuntimeError: 
Python builtin <built-in method random of Random object at 0x7fbcaaad5878> is currently not supported in Torchscript:
  File "/data/fangzhouhua/bluesky/icefall-master/egs/hik/asr1/pruned_transducer_stateless5/scaling.py", line 664
    def forward(self, x: Tensor) -> Tensor:
        if random.random() >= self.balance_prob:
           ~~~~~~~~~~~~~ <--- HERE
            return x
        else:

The ActivationBalancer, from what I understand, is to modify the backpropped derivatives and will not be used in inference mode? https://github.com/k2-fsa/icefall/blob/bd7fa2253dab9f627edc914b3289fb2f6c0e5bb6/egs/librispeech/ASR/pruned_transducer_stateless2/scaling.py#L652-L654

I think this could be fixed by returning input directly(when exporting) but Is there a way to fix this problem without touching scaling.py?

[csukuangfj]

which recipe are you using?

[csukuangfj]

Please update your code by following https://github.com/k2-fsa/icefall/blob/bd7fa2253dab9f627edc914b3289fb2f6c0e5bb6/egs/librispeech/ASR/pruned_transducer_stateless3/export.py#L664

which replaces ActivationBalancer with an Identity https://github.com/k2-fsa/icefall/blob/bd7fa2253dab9f627edc914b3289fb2f6c0e5bb6/egs/librispeech/ASR/pruned_transducer_stateless3/scaling_converter.py#L310-L312

[csukuangfj]

Also, would you mind creating a PR to fix the outdated code?

[Cescfangs]

Also, would you mind creating a PR to fix the outdated code?

Sure thing, I'll do it later.