add ema

toolbox/torchaudio/modules/utils/ema.py

@@ -66,6 +66,7 @@ class ErbEMA(nn.Module, EMANumpy):
         # erb_feat shape: [b, c, t, f]
         if state is None:
             state = self.make_erb_norm_state()
+            state = state.to(feat_erb.device)
         state = state.clone()
 
         for j in range(t):
@@ -118,6 +119,7 @@ class SpecEMA(nn.Module, EMANumpy):
         # feat_spec shape: [b, 2, t, df_bins]
         if state is None:
             state = self.make_spec_norm_state()
+            state = state.to(feat_spec.device)
         state = state.clone()
 
         for j in range(t):