TORCH_DEVICE_MODEL not TORCH_DEVICE

Author: iammosespaulr

surya/model/common/adetr/decoder.py

@@ -20,7 +20,7 @@
 
 
 def mark_step():
-    if settings.TORCH_DEVICE == 'xla':
+    if settings.TORCH_DEVICE_MODEL == 'xla':
         xm.mark_step()
 
 

surya/model/common/donut/encoder.py

@@ -21,7 +21,7 @@
 
 
 def mark_step():
-    if settings.TORCH_DEVICE == 'xla':
+    if settings.TORCH_DEVICE_MODEL == 'xla':
         xm.mark_step()
 
 
@@ -395,6 +395,8 @@ def forward(
         return outputs
 
 # Copied from transformers.models.swin.modeling_swin.SwinSelfOutput
+
+
 class DonutSwinSelfOutput(nn.Module):
     def __init__(self, config, dim):
         super().__init__()