diff --git a/extensions/machine-learning/src/views/models/prediction/columnsTable.ts b/extensions/machine-learning/src/views/models/prediction/columnsTable.ts
index d843147431..218fe6ab9e 100644
--- a/extensions/machine-learning/src/views/models/prediction/columnsTable.ts
+++ b/extensions/machine-learning/src/views/models/prediction/columnsTable.ts
@@ -326,7 +326,7 @@ export class ColumnsTable extends ModelViewBase implements IDataComponent<Predic
 				}
 
 				const currentColumn = columns.find(x => x.columnName === value);
-				if (currentColumn && modelParameter.type !== currentColumn?.dataType) {
+				if (currentColumn && modelParameter.type === currentColumn?.dataType) {
 					inputContainer.removeItem(warningButton);
 				} else {
 					inputContainer.addItem(warningButton, {