diff --git a/dvaas/packet_injection.cc b/dvaas/packet_injection.cc
index 57a78bc2..4e82aaf2 100644
--- a/dvaas/packet_injection.cc
+++ b/dvaas/packet_injection.cc
@@ -136,11 +136,6 @@ absl::StatusOr<PacketTestRuns> SendTestPacketsAndCollectOutputs(
             << packet_test_vector_by_id.size();
   statistics.total_packets_injected += packet_test_vector_by_id.size();
 
-  // Get IrP4Infos.
-  ASSIGN_OR_RETURN(const pdpi::IrP4Info sut_ir_p4info, GetIrP4Info(sut));
-  ASSIGN_OR_RETURN(const pdpi::IrP4Info control_ir_p4info,
-                   GetIrP4Info(control_switch));
-
   // Compute per packet injection delay.
   std::optional<absl::Duration> injection_delay;
   if (parameters.max_packets_to_send_per_second.has_value()) {
@@ -149,6 +144,8 @@ absl::StatusOr<PacketTestRuns> SendTestPacketsAndCollectOutputs(
   }
 
   // Send packets.
+  ASSIGN_OR_RETURN(const pdpi::IrP4Info control_ir_p4info,
+                   GetIrP4Info(control_switch));
   for (const auto& [test_id, packet_test_vector] : packet_test_vector_by_id) {
     if (packet_test_vector.input().type() == SwitchInput::DATAPLANE) {
       const Packet& packet = packet_test_vector.input().packet();
@@ -174,7 +171,7 @@ absl::StatusOr<PacketTestRuns> SendTestPacketsAndCollectOutputs(
   }
   LOG(INFO) << "Finished injecting test packets";
 
-  // Check the output of the switches.
+  // Check the output of the control switch.
   const absl::Duration kCollectionDuration = absl::Seconds(3);
   absl::StatusOr<std::vector<TaggedPacketIn>> control_packet_ins =
       CollectStreamMessageResponsesAndReturnTaggedPacketIns(
@@ -186,15 +183,6 @@ absl::StatusOr<PacketTestRuns> SendTestPacketsAndCollectOutputs(
             << " forwarded packets (from control switch)";
   statistics.total_packets_forwarded += control_packet_ins->size();
 
-  absl::StatusOr<std::vector<TaggedPacketIn>> sut_packet_ins =
-      CollectStreamMessageResponsesAndReturnTaggedPacketIns(
-          sut, kCollectionDuration, parameters.is_expected_unsolicited_packet);
-  RETURN_IF_ERROR(sut_packet_ins.status())
-      << "while  collecting the output of SUT";
-  LOG(INFO) << "Collected " << sut_packet_ins->size()
-            << " punted packets (from SUT)";
-  statistics.total_packets_punted += sut_packet_ins->size();
-
   absl::btree_map<int, SwitchOutput> switch_output_by_id;
   // Processing the output of the control switch.
   for (const TaggedPacketIn& packet_in : *control_packet_ins) {
@@ -217,21 +205,36 @@ absl::StatusOr<PacketTestRuns> SendTestPacketsAndCollectOutputs(
     *forwarded_output.mutable_port() = sut_egress_port.GetP4rtEncoding();
   }
 
-  // Processing the output of SUT.
-  for (const TaggedPacketIn& packet_in : *sut_packet_ins) {
-    // Add to (punted) switch output for ID.
-    PacketIn& punted_output =
-        *switch_output_by_id[packet_in.tag].add_packet_ins();
-
-    // Set hex and parsed packet.
-    punted_output.set_hex(
-        absl::BytesToHexString(packet_in.packet_in.payload()));
-    *punted_output.mutable_parsed() = packet_in.parsed_inner_packet;
-
-    // Set metadata.
-    ASSIGN_OR_RETURN(pdpi::IrPacketIn ir_packet_in,
-                     pdpi::PiPacketInToIr(sut_ir_p4info, packet_in.packet_in));
-    *punted_output.mutable_metadata() = ir_packet_in.metadata();
+  if (parameters.enable_sut_packet_in_collection) {
+    // Check the output of SUT.
+    absl::StatusOr<std::vector<TaggedPacketIn>> sut_packet_ins =
+        CollectStreamMessageResponsesAndReturnTaggedPacketIns(
+            sut, kCollectionDuration,
+            parameters.is_expected_unsolicited_packet);
+    RETURN_IF_ERROR(sut_packet_ins.status())
+        << "while  collecting the output of SUT";
+    LOG(INFO) << "Collected " << sut_packet_ins->size()
+              << " punted packets (from SUT)";
+    statistics.total_packets_punted += sut_packet_ins->size();
+
+    // Processing the output of SUT.
+    ASSIGN_OR_RETURN(const pdpi::IrP4Info sut_ir_p4info, GetIrP4Info(sut));
+    for (const TaggedPacketIn& packet_in : *sut_packet_ins) {
+      // Add to (punted) switch output for ID.
+      PacketIn& punted_output =
+          *switch_output_by_id[packet_in.tag].add_packet_ins();
+
+      // Set hex and parsed packet.
+      punted_output.set_hex(
+          absl::BytesToHexString(packet_in.packet_in.payload()));
+      *punted_output.mutable_parsed() = packet_in.parsed_inner_packet;
+
+      // Set metadata.
+      ASSIGN_OR_RETURN(
+          pdpi::IrPacketIn ir_packet_in,
+          pdpi::PiPacketInToIr(sut_ir_p4info, packet_in.packet_in));
+      *punted_output.mutable_metadata() = ir_packet_in.metadata();
+    }
   }
 
   // Create PacketTestRuns.
diff --git a/dvaas/packet_injection.h b/dvaas/packet_injection.h
index 3c398be5..f2c117bb 100644
--- a/dvaas/packet_injection.h
+++ b/dvaas/packet_injection.h
@@ -83,6 +83,12 @@ struct PacketInjectionParams {
   // The mapping of P4RT port IDs for connected interfaces between SUT and the
   // control switch.
   MirrorTestbedP4rtPortIdMap mirror_testbed_port_map;
+  // If false, does not collect packet-ins from SUT. Useful for scenarios where
+  // the connection to SUT is expected to go down such as in NSF and we don't
+  // want the injection/collection function to fail because of that.
+  // TODO: Replace with parameter to tolerate SUT disconnections
+  // rather than ignoring all packet-ins.
+  bool enable_sut_packet_in_collection = true;
 };
 
 // Determines the switch's behavior when receiving test packets by:
diff --git a/dvaas/traffic_generator.cc b/dvaas/traffic_generator.cc
index ccb3118c..137ef230 100644
--- a/dvaas/traffic_generator.cc
+++ b/dvaas/traffic_generator.cc
@@ -160,6 +160,9 @@ void SimpleTrafficGenerator::InjectTraffic() {
             .mirror_testbed_port_map =
                 params_.validation_params.mirror_testbed_port_map_override
                     .value_or(MirrorTestbedP4rtPortIdMap::CreateIdentityMap()),
+            .enable_sut_packet_in_collection =
+                !params_.validation_params.switch_output_diff_params
+                     .treat_expected_and_actual_outputs_as_having_no_packet_ins,
         },
         statistics);
     CHECK_OK(test_runs.status());  // Crash OK.