diff --git a/src/wifi/model/nist-error-rate-model.cc b/src/wifi/model/nist-error-rate-model.cc
index a62fbf36f..2e374f161 100644
--- a/src/wifi/model/nist-error-rate-model.cc
+++ b/src/wifi/model/nist-error-rate-model.cc
@@ -275,7 +275,7 @@ NistErrorRateModel::GetFec1024QamBer (double snr, uint64_t nbits,
 double
 NistErrorRateModel::DoGetChunkSuccessRate (WifiMode mode, WifiTxVector txVector, double snr, uint64_t nbits) const
 {
-  NS_LOG_FUNCTION (this << mode << txVector.GetMode () << snr << nbits);
+  NS_LOG_FUNCTION (this << mode << snr << nbits);
   if (mode.GetModulationClass () == WIFI_MOD_CLASS_ERP_OFDM
       || mode.GetModulationClass () == WIFI_MOD_CLASS_OFDM
       || mode.GetModulationClass () == WIFI_MOD_CLASS_HT
diff --git a/src/wifi/model/wifi-phy-state-helper.cc b/src/wifi/model/wifi-phy-state-helper.cc
index 3bab853af..a606ef072 100644
--- a/src/wifi/model/wifi-phy-state-helper.cc
+++ b/src/wifi/model/wifi-phy-state-helper.cc
@@ -475,13 +475,13 @@ WifiPhyStateHelper::ContinueRxNextMpdu (Ptr<WifiPsdu> psdu, double snr, WifiTxVe
 }
 
 void
-WifiPhyStateHelper::SwitchFromRxEndOk (Ptr<WifiPsdu> psdu, double snr, WifiTxVector txVector, std::vector<bool> statusPerMpdu)
+WifiPhyStateHelper::SwitchFromRxEndOk (Ptr<WifiPsdu> psdu, double snr, WifiTxVector txVector, uint16_t staId, std::vector<bool> statusPerMpdu)
 {
   NS_LOG_FUNCTION (this << *psdu << snr << txVector << statusPerMpdu.size () <<
                    std::all_of(statusPerMpdu.begin(), statusPerMpdu.end(), [](bool v) { return v; })); //returns true if all true
   NS_ASSERT (statusPerMpdu.size () != 0);
   NS_ASSERT (m_endRx == Simulator::Now ());
-  m_rxOkTrace (psdu->GetPacket (), snr, txVector.GetMode (), txVector.GetPreambleType ());
+  m_rxOkTrace (psdu->GetPacket (), snr, txVector.GetMode (staId), txVector.GetPreambleType ());
   NotifyRxEndOk ();
   DoSwitchFromRx ();
   if (!m_rxOkCallback.IsNull ())
diff --git a/src/wifi/model/wifi-phy-state-helper.h b/src/wifi/model/wifi-phy-state-helper.h
index 142d0cf95..14fac90bc 100644
--- a/src/wifi/model/wifi-phy-state-helper.h
+++ b/src/wifi/model/wifi-phy-state-helper.h
@@ -197,9 +197,10 @@ public:
    * \param psdu the successfully received PSDU
    * \param snr the SNR of the received PSDU in linear scale
    * \param txVector TXVECTOR of the PSDU
+   * \param staId the station ID of the PSDU (only used for MU)
    * \param statusPerMpdu reception status per MPDU
    */
-  void SwitchFromRxEndOk (Ptr<WifiPsdu> psdu, double snr, WifiTxVector txVector, std::vector<bool> statusPerMpdu);
+  void SwitchFromRxEndOk (Ptr<WifiPsdu> psdu, double snr, WifiTxVector txVector, uint16_t staId, std::vector<bool> statusPerMpdu);
   /**
    * Switch from RX after the reception failed.
    *
diff --git a/src/wifi/model/wifi-phy.cc b/src/wifi/model/wifi-phy.cc
index 5742f73b2..7079b8099 100644
--- a/src/wifi/model/wifi-phy.cc
+++ b/src/wifi/model/wifi-phy.cc
@@ -3050,7 +3050,7 @@ WifiPhy::StartReceivePayload (Ptr<Event> event)
       Ptr<const WifiPsdu> psdu = GetAddressedPsduInPpdu (ppdu);
       if (psdu)
         {
-          WifiMode txMode = txVector.GetPreambleType () == WIFI_PREAMBLE_HE_MU ? txVector.GetMode (GetStaId ()) : txVector.GetMode ();
+          WifiMode txMode = txVector.GetMode (GetStaId ());
           uint8_t nss = txVector.GetNssMax();
           if (txVector.GetPreambleType () == WIFI_PREAMBLE_HE_MU)
             {
@@ -3205,7 +3205,7 @@ WifiPhy::EndReceive (Ptr<Event> event)
       //At least one MPDU has been successfully received
       WifiTxVector txVector = event->GetTxVector ();
       NotifyMonitorSniffRx (psdu, GetFrequency (), txVector, m_signalNoise, m_statusPerMpdu);
-      m_state->SwitchFromRxEndOk (Copy (psdu), snr, txVector, m_statusPerMpdu);
+      m_state->SwitchFromRxEndOk (Copy (psdu), snr, txVector, staId, m_statusPerMpdu);
     }
   else
     {
diff --git a/src/wifi/model/wifi-tx-vector.cc b/src/wifi/model/wifi-tx-vector.cc
index cdea8d177..061a85897 100644
--- a/src/wifi/model/wifi-tx-vector.cc
+++ b/src/wifi/model/wifi-tx-vector.cc
@@ -377,13 +377,11 @@ WifiTxVector::GetHeMuUserInfoMap (void) const
 
 std::ostream & operator << ( std::ostream &os, const WifiTxVector &v)
 {
-  os << "mode: " << v.GetMode ()
-     << " txpwrlvl: " << +v.GetTxPowerLevel ()
+  os << "txpwrlvl: " << +v.GetTxPowerLevel ()
      << " preamble: " << v.GetPreambleType ()
      << " channel width: " << v.GetChannelWidth ()
      << " GI: " << v.GetGuardInterval ()
      << " NTx: " << +v.GetNTx ()
-     << " Nss: " << +v.GetNss ()
      << " Ness: " << +v.GetNess ()
      << " MPDU aggregation: " << v.IsAggregation ()
      << " STBC: " << v.IsStbc ()
@@ -400,6 +398,11 @@ std::ostream & operator << ( std::ostream &os, const WifiTxVector &v)
              << ", Nss: " << +ui.second.nss << "}";
         }
     }
+  else
+    {
+      os << " mode: " << v.GetMode ()
+         << " Nss: " << +v.GetNss ();
+    }
   return os;
 }