Merge pull request #45 from PathwayCommons/development

Use mixed precision
PathwayCommons · May 26, 2022 · f033f9d · f033f9d
2 parents b18897f + 08a64ea
commit f033f9d
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/pathway_abstract_classifier/pathway_abstract_classifier.py b/pathway_abstract_classifier/pathway_abstract_classifier.py
@@ -3,6 +3,12 @@
 import ktrain
 from cached_path import cached_path
 
+from tensorflow.keras import mixed_precision  # erroneous missing import
+import tensorflow as tf
+
+if len(tf.config.list_physical_devices('GPU')) > 0:
+    mixed_precision.set_global_policy('mixed_float16')
+
 
 class Prediction(NamedTuple):
     document: Dict[str, str]