More fixes for issue
@@ -204,6 +204,8 @@ class ParserModel(Model):
if new_output == self.upper.nO:
return
smaller = self.upper
+
+ with Model.use_device('cpu'):
larger = Affine(new_output, smaller.nI)
# Set nan as value for unseen classes, to prevent prediction.
larger.W.fill(self.ops.xp.nan)
