fix(hook): fix torch.Tensor support in ArtifactHook

pyannote · May 8, 2024 · 4615808 · 4615808
1 parent 7a90137
commit 4615808
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -16,6 +16,7 @@
 
 - fix(task): fix random generators and their reproducibility (with [@FrenchKrab](https://github.com/FrenchKrab))
 - fix(task): fix estimation of training set size (with [@FrenchKrab](https://github.com/FrenchKrab))
+- fix(hook): fix `torch.Tensor` support in `ArtifactHook`
 
 ### Improvements
 

diff --git a/pyannote/audio/pipelines/utils/hook.py b/pyannote/audio/pipelines/utils/hook.py
@@ -24,6 +24,7 @@
 from copy import deepcopy
 from typing import Any, Mapping, Optional, Text
 
+import torch
 from rich.progress import (
     BarColumn,
     Progress,
@@ -75,6 +76,9 @@ def __call__(
         ):
             return
 
+        if isinstance(step_artifact, torch.Tensor):
+            step_artifact = step_artifact.numpy(force=True)
+
         file.setdefault(self.file_key, dict())[step_name] = deepcopy(step_artifact)