Document FP16 inference

README.md

@@ -20,28 +20,38 @@ pip install structured-data-models
 Use labeled examples as context to predict class probabilities for new data:
 
 ```python
+import torch
 from sklearn.datasets import load_breast_cancer
+
 import sdm
 
 df = load_breast_cancer(as_frame=True).frame
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 table = sdm.TableTensor.from_pandas(
     df=df,
     stypes=sdm.infer_stypes(df, overrides={"target": "categorical"}),
-    device="cuda",
-)
-model = sdm.models.KumoTabular(task="classification", device="cuda")
-
-probs = model(
-    x_context=table[:300].drop_columns("target"),
-    y_context=table[:300, "target"],
-    x_query=table[300:].drop_columns("target"),
-    num_estimators=8,
+    device=device,
 )
+model = sdm.models.KumoTabular(task="classification", device=device)
+
+with torch.amp.autocast(
+    device.type,
+    dtype=torch.float16,
+    enabled=device.type == "cuda",
+):
+    probs = model(
+        x_context=table[:300].drop_columns("target"),
+        y_context=table[:300, "target"],
+        x_query=table[300:].drop_columns("target"),
+        num_estimators=8,
+    )
 
 print(probs)
 ```
 
+Published GPU benchmarks use FP16 mixed precision for model execution. SDM respects the caller's active PyTorch autocast context but does not enable autocasting automatically, so use the context shown above for comparable GPU throughput.
+
 To learn more, visit [structured-data-models](https://github.com/NVIDIA/structured-data-models).
 
 ## License