Update think_sound/models/mmmodules/model/embeddings.py

think_sound/models/mmmodules/model/embeddings.py

@@ -21,7 +21,7 @@ class TimestepEmbedder(nn.Module):
         assert dim % 2 == 0, 'dim must be even.'
 
         with torch.autocast('cuda', enabled=False):
-            self.freqs = nn.Buffer(
+            self.register_buffer("freqs",
                 1.0 / (10000**(torch.arange(0, frequency_embedding_size, 2, dtype=torch.float32) /
                                frequency_embedding_size)),
                 persistent=False)