Fix issue with bf16

AmpereComputingAI · Sep 25, 2023 · e6cc34c · e6cc34c
1 parent dd0d2d0
commit e6cc34c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/computer_vision/classification/vgg_16/run.py b/computer_vision/classification/vgg_16/run.py
@@ -91,7 +91,7 @@ def run_pytorch_fp(model_name, batch_size, num_runs, timeout, images_path, label
 
     def run_single_pass(pytorch_runner, imagenet):
         shape = (224, 224)
-        output = pytorch_runner.run(batch_size, torch.from_numpy(imagenet.get_input_array(shape)))
+        output = pytorch_runner.run(batch_size, torch.from_numpy(imagenet.get_input_array(shape))).float()
 
         for i in range(batch_size):
             imagenet.submit_predictions(