Fix bfloat16 numpy conversion

cinema/examples/inference/classification_cvd.py

@@ -50,9 +50,9 @@ def run(trained_dataset: str, view: str, seed: int, device: torch.device, dtype:
     with torch.no_grad(), torch.autocast("cuda", dtype=dtype, enabled=torch.cuda.is_available()):
         logits = model(batch)  # (1, n_classes)
     probs = torch.softmax(logits, dim=1)[0]  # (n_classes,)
-    probs_dict = dict(zip(classes, probs.cpu().numpy(), strict=False))
+    probs_dict = dict(zip(classes, probs.cpu().float().numpy(), strict=False))
     print(f"Using {view} view with model trained on {trained_dataset} dataset with seed {seed}.")  # noqa: T201
-    print(f"The predicted class is {classes[np.argmax(logits)]}, ground truth should be HCM.")  # noqa: T201
+    print(f"The predicted class is {classes[np.argmax(logits.cpu().float().numpy())]}, ground truth should be HCM.")  # noqa: T201
     print(f"The probabilities are {probs_dict}.")  # noqa: T201
 
 

cinema/examples/inference/classification_sex.py

@@ -51,9 +51,9 @@ def run(seed: int, device: torch.device, dtype: torch.dtype) -> None:
     with torch.no_grad(), torch.autocast("cuda", dtype=dtype, enabled=torch.cuda.is_available()):
         logits = model(batch)  # (1, n_classes)
     probs = torch.softmax(logits, dim=1)[0]  # (n_classes,)
-    probs_dict = dict(zip(classes, probs.cpu().numpy(), strict=False))
+    probs_dict = dict(zip(classes, probs.cpu().float().numpy(), strict=False))
     print(f"Using {view} view with model trained on {trained_dataset} dataset with seed {seed}.")  # noqa: T201
-    print(f"The predicted class is {classes[np.argmax(logits)]}, ground truth should be Female.")  # noqa: T201
+    print(f"The predicted class is {classes[np.argmax(logits.cpu().float().numpy())]}, ground truth should be Female.")  # noqa: T201
     print(f"The probabilities are {probs_dict}.")  # noqa: T201
 
 

cinema/examples/inference/classification_vendor.py

@@ -51,9 +51,9 @@ def run(view: str, seed: int, device: torch.device, dtype: torch.dtype) -> None:
     with torch.no_grad(), torch.autocast("cuda", dtype=dtype, enabled=torch.cuda.is_available()):
         logits = model(batch)  # (1, n_classes)
     probs = torch.softmax(logits, dim=1)[0]  # (n_classes,)
-    probs_dict = dict(zip(classes, probs.cpu().numpy(), strict=False))
+    probs_dict = dict(zip(classes, probs.cpu().float().numpy(), strict=False))
     print(f"Using {view} view with model trained on {trained_dataset} dataset with seed {seed}.")  # noqa: T201
-    print(f"The predicted class is {classes[np.argmax(logits)]}, ground truth should be SIEMENS.")  # noqa: T201
+    print(f"The predicted class is {classes[np.argmax(logits.cpu().float().numpy())]}, ground truth should be SIEMENS.")  # noqa: T201
     print(f"The probabilities are {probs_dict}.")  # noqa: T201
 
 

cinema/examples/inference/landmark_coordinate.py

@@ -35,7 +35,7 @@ def run(view: str, seed: int, device: torch.device, dtype: torch.dtype) -> None:
         batch = transform({view: torch.from_numpy(images[None, ..., 0, t])})
         batch = {k: v[None, ...].to(device=device, dtype=dtype) for k, v in batch.items()}
         with torch.no_grad(), torch.autocast("cuda", dtype=dtype, enabled=torch.cuda.is_available()):
-            coords = model(batch)[0].numpy()  # (6,)
+            coords = model(batch)[0].cpu().float().numpy()  # (6,)
         coords *= np.array([w, h, w, h, w, h])
         coords = [int(x) for x in coords]
         coords_list.append(coords)

cinema/examples/inference/landmark_heatmap.py

@@ -214,7 +214,7 @@ def run(view: str, seed: int, device: torch.device, dtype: torch.dtype) -> None:
             logits = model(batch)[view]  # (1, 3, x, y)
         probs = torch.sigmoid(logits)  # (1, 3, width, height)
         probs_list.append(probs[0].detach().to(torch.float32).cpu().numpy())
-        coords = heatmap_soft_argmax(probs)[0].numpy()
+        coords = heatmap_soft_argmax(probs)[0].cpu().float().numpy()
         coords = [int(x) for x in coords]
         coords_list.append(coords)
     probs = np.stack(probs_list, axis=-1)  # (3, x, y, t)