Force SGD to use momentums to bypass an MLIR bridge issue.
@@ -138,7 +138,6 @@ class SGD(optimizer.Optimizer):
if hasattr(self, "_built") and self._built:
return
self.momentums = []
- if self.momentum != 0:
for var in var_list:
self.momentums.append(
self.add_variable_from_reference(
@@ -152,7 +151,6 @@ class SGD(optimizer.Optimizer):
lr = tf.cast(self.learning_rate, variable.dtype)
m = None
var_key = self._var_key(variable)
- if self.momentum != 0:
momentum = tf.cast(self.momentum, variable.dtype)
m = self.momentums[self._index_dict[var_key]]
