zchen/b2txt25
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

tpu

Browse Source
This commit is contained in:
Zchen
2025-10-15 00:18:05 +08:00
parent 4a3d3f35ec
commit 603bb12220
1 changed files with 16 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
model_training_nnn/rnn_model.py
View File

@@ -117,7 +117,9 @@ class NoiseModel(nn.Module):
if states.dtype != gru_dtype:
states = states.to(gru_dtype)
# GRU forward pass
# Disable autocast for GRU to avoid dtype mismatches on XLA
device_type = x.device.type
with torch.autocast(device_type=device_type, enabled=False):
output, hidden_states = self.gru(x, states)
return output, hidden_states
@@ -225,7 +227,8 @@ class CleanSpeechModel(nn.Module):
if states.dtype != gru_dtype:
states = states.to(gru_dtype)
# GRU forward pass
device_type = x.device.type
with torch.autocast(device_type=device_type, enabled=False):
output, hidden_states = self.gru(x, states)
# Classification
@@ -309,7 +312,8 @@ class NoisySpeechModel(nn.Module):
if states.dtype != gru_dtype:
states = states.to(gru_dtype)
# GRU forward pass
device_type = x.device.type
with torch.autocast(device_type=device_type, enabled=False):
output, hidden_states = self.gru(x, states)
# Classification
@@ -444,6 +448,8 @@ class TripleGRUDecoder(nn.Module):
states = states.to(clean_gru_dtype)
# GRU forward pass (skip preprocessing since input is already processed)
device_type = x_processed.device.type
with torch.autocast(device_type=device_type, enabled=False):
output, hidden_states = self.clean_speech_model.gru(x_processed, states)
# Classification
@@ -466,6 +472,8 @@ class TripleGRUDecoder(nn.Module):
states = states.to(noisy_gru_dtype)
# GRU forward pass (NoisySpeechModel doesn't have day layers anyway)
device_type = x_processed.device.type
with torch.autocast(device_type=device_type, enabled=False):
output, hidden_states = self.noisy_speech_model.gru(x_processed, states)
# Classification