Robustness is a key requirement for deploying a machine learning (ML) based solution. When a solution involves a ML model whose robustness is not guaranteed, ensuring robustness of the solution might rely on continuous checking of the ML model for its validity after the solution is deployed in production. Using wafer image classification as an example, this paper introduces tensor-based methods that help improve robustness of a neural-network-based classification approach and facilitate its deployment. Experiment results based on data from a commercial product line are presented to explain the key ideas behind the tensor-based methods.