dtype can be selectable

__main__.py

@@ -16,6 +16,7 @@ parser.add_argument("references", type=str, help="Path to the numpy array contai
 parser.add_argument("--output", type=str, help="Path to the output file")
 parser.add_argument("--batch_size", type=int, help="Batch size to use for the computation")
 parser.add_argument("--num_process", type=int, help="Batch size to use for the computation", default=1)
+parser.add_argument("--dtype", type=str, help="Data type to use for the computation", default="float32")
 parser.add_argument("--debug", action="store_true", help="Debug mode")
 args = parser.parse_args()
 
@@ -23,8 +24,8 @@ if not args.predictions or not args.references:
     raise ValueError("You must provide the path to the predictions and references numpy arrays")
 
 
-predictions = np.load(args.predictions).astype(np.float32)
-references = np.load(args.references).astype(np.float32)
+predictions = np.load(args.predictions).astype(args.dtype)
+references = np.load(args.references).astype(args.dtype)
 
 if args.debug:
     predictions = predictions[:1000]
@@ -45,6 +46,7 @@ results = metric.compute(
     num_process=args.num_process,
     return_each_features=True,
     return_coverages=True,
+    dtype=args.dtype,
 )
 logger.info(f"Time taken: {time.time() - s}")