diff --git a/src/wifi/model/wifi-tx-vector.cc b/src/wifi/model/wifi-tx-vector.cc
index e2be9b5a6..cdea8d177 100644
--- a/src/wifi/model/wifi-tx-vector.cc
+++ b/src/wifi/model/wifi-tx-vector.cc
@@ -255,6 +255,14 @@ WifiTxVector::SetNss (uint8_t nss)
   m_nss = nss;
 }
 
+void
+WifiTxVector::SetNss (uint8_t nss, uint16_t staId)
+{
+  NS_ABORT_MSG_IF (m_preamble != WIFI_PREAMBLE_HE_MU, "Not an HE MU transmission");
+  NS_ABORT_MSG_IF (staId > 2048, "STA-ID should be correctly set for HE MU");
+  m_muUserInfos[staId].nss = nss;
+}
+
 void
 WifiTxVector::SetNess (uint8_t ness)
 {
diff --git a/src/wifi/model/wifi-tx-vector.h b/src/wifi/model/wifi-tx-vector.h
index d8f051430..5365f8be4 100644
--- a/src/wifi/model/wifi-tx-vector.h
+++ b/src/wifi/model/wifi-tx-vector.h
@@ -202,17 +202,24 @@ public:
    */
   uint8_t GetNssMax (void) const;
   /**
-   * Sets the number of Nss refer to IEEE 802.11n Table 20-28 for explanation and range
+   * Sets the number of Nss
    *
    * \param nss the number of spatial streams
    */
   void SetNss (uint8_t nss);
+  /**
+   * Sets the number of Nss for HE MU
+   *
+   * \param nss the number of spatial streams
+   * \param staId the station ID for HE MU
+   */
+  void SetNss (uint8_t nss, uint16_t staId);
   /**
    * \returns the number of extended spatial streams
    */
   uint8_t GetNess (void) const;
   /**
-   * Sets the Ness number refer to IEEE 802.11n Table 20-6 for explanation
+   * Sets the Ness number
    *
    * \param ness the number of extended spatial streams
    */