format

torchrl/modules/tensordict_module/rnn.py

@@ -127,6 +127,8 @@ def lstm_cell(self, x, hx, cx):
 
 # copy LSTM
 class LSTMBase(nn.RNNBase):
+    """A Base module for LSTM. Inheriting from LSTMBase enables compatibility with torch.compile."""
+
     def __init__(self, *args, **kwargs):
         return super().__init__("LSTM", *args, **kwargs)
 
@@ -137,8 +139,7 @@ def __init__(self, *args, **kwargs):
 
 
 class LSTM(LSTMBase):
-    """A PyTorch module for executing multiple steps of a multi-layer LSTM. The module behaves exactly like
-    :class:`torch.nn.LSTM`, but this implementation is exclusively coded in Python.
+    """A PyTorch module for executing multiple steps of a multi-layer LSTM. The module behaves exactly like :class:`torch.nn.LSTM`, but this implementation is exclusively coded in Python.
 
     .. note::
         This class is implemented without relying on CuDNN, which makes it compatible with :func:`torch.vmap` and :func:`torch.compile`.
@@ -811,6 +812,8 @@ def gru_cell(self, x, hx):
 
 # copy GRU
 class GRUBase(nn.RNNBase):
+    """A Base module for GRU. Inheriting from GRUBase enables compatibility with torch.compile."""
+
     def __init__(self, *args, **kwargs):
         return super().__init__("GRU", *args, **kwargs)
 
@@ -821,8 +824,7 @@ def __init__(self, *args, **kwargs):
 
 
 class GRU(GRUBase):
-    """A PyTorch module for executing multiple steps of a multi-layer GRU. The module behaves exactly like
-    :class:`torch.nn.GRU`, but this implementation is exclusively coded in Python.
+    """A PyTorch module for executing multiple steps of a multi-layer GRU. The module behaves exactly like :class:`torch.nn.GRU`, but this implementation is exclusively coded in Python.
 
     .. note::
         This class is implemented without relying on CuDNN, which makes it compatible with :func:`torch.vmap` and :func:`torch.compile`.