plot float16 alignments

2026-09-09 11:16:00 +08:00 · 2021-01-06 13:19:40 +01:00
parent 13c6665c92
commit e4680e1b99
1 changed files with 2 additions and 0 deletions
@@ -17,6 +17,8 @@ def plot_alignment(alignment,
        alignment_ = alignment.detach().cpu().numpy().squeeze()
    else:
        alignment_ = alignment
+    alignment_ = alignment_.astype(
+        np.float32) if alignment_.dtype == np.float16 else alignment_
    fig, ax = plt.subplots(figsize=fig_size)
    im = ax.imshow(alignment_.T,
                   aspect='auto',