diff --git a/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc b/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc
index be8aeb3e6..208501758 100644
--- a/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc
+++ b/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include <stdlib.h>
 
+#include <limits>
 #include <string>
 
 #include "../qsim/lib/channel.h"
@@ -60,7 +61,10 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {
 
   void Compute(tensorflow::OpKernelContext* context) override {
     // TODO (mbbrough): add more dimension checks for other inputs here.
-    DCHECK_EQ(4, context->num_inputs());
+    const int num_inputs = context->num_inputs();
+    OP_REQUIRES(context, num_inputs == 4,
+                tensorflow::errors::InvalidArgument(absl::StrCat(
+                    "Expected 4 inputs, got ", num_inputs, " inputs.")));
 
     // Parse to Program Proto and num_qubits.
     std::vector<Program> programs;
@@ -79,6 +83,12 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {
 
     int num_samples = 0;
     OP_REQUIRES_OK(context, GetIndividualSample(context, &num_samples));
+    OP_REQUIRES(
+        context,
+        num_samples >= 0 && num_samples < std::numeric_limits<int>::max(),
+        tensorflow::errors::InvalidArgument(
+            absl::StrCat("num_samples must be between 0 and ",
+                         std::numeric_limits<int>::max(), ".")));
 
     // Construct qsim circuits.
     std::vector<NoisyQsimCircuit> qsim_circuits(programs.size(),