From 72ca831dd483a930841be7ab99b87f08ef83cde7 Mon Sep 17 00:00:00 2001
From: Ross Wightman <rwightman@gmail.com>
Date: Wed, 19 May 2021 12:01:53 -0700
Subject: [PATCH] Back to using strings for the enum translation, forgot about
 import dep

---
 timm/bits/device_env_xla.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/timm/bits/device_env_xla.py b/timm/bits/device_env_xla.py
index 71d350fd..a565c1c8 100644
--- a/timm/bits/device_env_xla.py
+++ b/timm/bits/device_env_xla.py
@@ -23,12 +23,12 @@ from .device_env import DeviceEnv, DeviceEnvType, TensorList
 
 
 _PT_TO_XM_OP = {
-    ReduceOp.SUM: xm.REDUCE_SUM,
-    ReduceOp.PRODUCT: xm.REDUCE_MUL,
-    ReduceOp.MIN: xm.REDUCE_MIN,
-    ReduceOp.MAX:  xm.REDUCE_MAX,
-    ReduceOp.BAND: xm.REDUCE_AND,
-    ReduceOp.BOR: xm.REDUCE_OR,
+    ReduceOp.SUM: 'sum',
+    ReduceOp.PRODUCT: 'mul',
+    ReduceOp.MIN: 'min',
+    ReduceOp.MAX: 'max',
+    ReduceOp.BAND: 'and',
+    ReduceOp.BOR: 'or',
 }