From 4b08a7e39c16567b5efaf456dde7554cf525abc1 Mon Sep 17 00:00:00 2001
From: Andrei Panferov <andrei@panferov.org>
Date: Sun, 18 Feb 2024 12:55:51 +0100
Subject: [PATCH] comment

---
 inference_lib/src/aqlm/inference_kernels/kernel_selector.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/inference_lib/src/aqlm/inference_kernels/kernel_selector.py b/inference_lib/src/aqlm/inference_kernels/kernel_selector.py
index 2994aa3a..bb1a55df 100644
--- a/inference_lib/src/aqlm/inference_kernels/kernel_selector.py
+++ b/inference_lib/src/aqlm/inference_kernels/kernel_selector.py
@@ -12,6 +12,7 @@
 
 @contextmanager
 def optimize_for_training():
+    """Use this context manager during model initialization (e.g. `.from_pretrained(...)`) to select inference kernels optimized for larger batch sizes"""
     global _OPTIMIZE_FOR_TRAINING
     _OPTIMIZE_FOR_TRAINING = True
     try: