diff --git a/models/Model_XSeg/Model.py b/models/Model_XSeg/Model.py
index fe444d2..4ab23fc 100644
--- a/models/Model_XSeg/Model.py
+++ b/models/Model_XSeg/Model.py
@@ -101,8 +101,8 @@ class XSegModel(ModelBase):
 
 
             # Average losses and gradients, and create optimizer update ops
-            with tf.device(f'/CPU:0'): # Temporary fix. Unknown bug with training freeze starts from 2.4.0, but 2.3.1 was ok
-            #with tf.device (models_opt_device):
+            #with tf.device(f'/CPU:0'): # Temporary fix. Unknown bug with training freeze starts from 2.4.0, but 2.3.1 was ok
+            with tf.device (models_opt_device):
                 pred = tf.concat(gpu_pred_list, 0)
                 loss = tf.concat(gpu_losses, 0)
                 loss_gv_op = self.model.opt.get_update_op (nn.average_gv_list (gpu_loss_gvs))