Skip to content

Commit

Permalink
Fix an import in export_utils.py (#8571)
Browse files Browse the repository at this point in the history
Signed-off-by: w4-jinhyeonkim <131935801+w4-jinhyeonkim@users.noreply.github.com>
Signed-off-by: Pablo Garay <pagaray@nvidia.com>
  • Loading branch information
w4-jinhyeonkim authored and pablo-garay committed Mar 19, 2024
1 parent ee3660a commit ae752d1
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions nemo/utils/export_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -448,8 +448,6 @@ def replace_for_export(model: nn.Module) -> nn.Module:
Returns:
model, possibly modified in-place
"""
from nemo.collections.tts.modules.submodules import MaskedInstanceNorm1d

default_replacements = {
"MatchedScaleMaskSoftmax": wrap_module(None, replace_MatchedScaleMaskSoftmax),
}
Expand All @@ -466,6 +464,8 @@ def add_casts_around_norms(model: nn.Module):
It was used with an extra post-parse script to have TRT preserve extra precision when --fp16 needed.
Should not be needed with TRT 8.6.1 or later.
"""
from nemo.collections.tts.modules.submodules import MaskedInstanceNorm1d

default_cast_replacements = {
"BatchNorm1d": wrap_module(nn.BatchNorm1d, CastToFloat),
"BatchNorm2d": wrap_module(nn.BatchNorm2d, CastToFloat),
Expand Down

0 comments on commit ae752d1

Please sign in to comment.