diff options
Diffstat (limited to 'swr2_asr/utils/visualization.py')
| -rw-r--r-- | swr2_asr/utils/visualization.py | 4 |
1 files changed, 3 insertions, 1 deletions
diff --git a/swr2_asr/utils/visualization.py b/swr2_asr/utils/visualization.py index b288c5a..23956fd 100644 --- a/swr2_asr/utils/visualization.py +++ b/swr2_asr/utils/visualization.py @@ -14,7 +14,9 @@ def plot(path): epoch = 5 while True: try: - current_state = torch.load(path + str(epoch), map_location=torch.device("cpu")) + current_state = torch.load( + path + str(epoch), map_location=torch.device("cpu") + ) # pylint: disable=no-member except FileNotFoundError: break train_losses.append((epoch, current_state["train_loss"].item())) |