inference mode doesn't work with torch xla

Author: iammosespaulr

surya/detection.py

@@ -88,7 +88,7 @@ def batch_detection(
         if static_cache:
             batch = pad_to_batch_size(batch, batch_size)
 
-        with torch.inference_mode():
+        with torch.no_grad():
             pred = model(pixel_values=batch)
 
         logits = pred.logits

surya/layout.py

@@ -116,7 +116,7 @@ def batch_layout_detection(images: List, model, processor, batch_size=None) -> L
 
         batch_predictions = [[] for _ in range(current_batch_size)]
 
-        with torch.inference_mode():
+        with torch.no_grad():
             encoder_hidden_states = model.encoder(pixel_values=batch_pixel_values)[0]
 
             token_count = 0

surya/recognition.py

@@ -88,7 +88,7 @@ def batch_recognition(images: List[Image.Image], languages: List[List[str] | Non
         all_done = torch.zeros(current_batch_size, dtype=torch.bool, device=model.device)
         encoder_hidden_states = None
 
-        with torch.inference_mode():
+        with torch.no_grad():
             encoder_batch_size = batch_size // settings.RECOGNITION_ENCODER_BATCH_DIVISOR
             for z in range(0, batch_pixel_values.shape[0], encoder_batch_size):
                 encoder_pixel_values = batch_pixel_values[z:min(z + encoder_batch_size, batch_pixel_values.shape[0])]

surya/tables.py

@@ -100,7 +100,7 @@ def batch_table_recognition(images: List, table_cells: List[List[Dict]], model:
 
         batch_predictions = [[] for _ in range(current_batch_size)]
 
-        with torch.inference_mode():
+        with torch.no_grad():
             encoder_hidden_states = model.encoder(pixel_values=batch_pixel_values).last_hidden_state
             text_encoder_hidden_states = model.text_encoder(
                 input_boxes=batch_bboxes,