YoussefMoNader commited on
Commit
6ebe1fe
·
verified ·
1 Parent(s): f555450

Update ink_detection_pipeline.py

Browse files

fixing bfloat16 handling on certain devices

Files changed (1) hide show
  1. ink_detection_pipeline.py +1 -1
ink_detection_pipeline.py CHANGED
@@ -72,7 +72,7 @@ class InkDetectionPipeline(Pipeline):
72
  sub_y_preds = torch.sigmoid(sub_y_preds)
73
 
74
  # Move to CPU and numpy
75
- sub_y_preds = sub_y_preds.detach().cpu().numpy()
76
  # shape (subB, 1, tile_size, tile_size)
77
 
78
  all_preds.append(sub_y_preds)
 
72
  sub_y_preds = torch.sigmoid(sub_y_preds)
73
 
74
  # Move to CPU and numpy
75
+ sub_y_preds = sub_y_preds.detach().cpu().float().numpy()
76
  # shape (subB, 1, tile_size, tile_size)
77
 
78
  all_preds.append(sub_y_preds)