fixed batch first bug

stefanonardo · Apr 12, 2018 · a053adc · a053adc
1 parent 259a261
commit a053adc
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 3 deletions.
diff --git a/README.md b/README.md
@@ -10,7 +10,7 @@ PyTorch-ESN is a PyTorch module implementing Echo State Networks with leaky-inte
 
 ### Offline training (ridge regression)
 
-```
+```python
 from torchesn.nn import ESN
 from torchesn.utils import prepare_target
 

diff --git a/setup.py b/setup.py
@@ -1,7 +1,7 @@
 from setuptools import setup, find_packages
 
 setup(name='pytorch-esn',
-      version='1.0.2',
+      version='1.0.3',
       packages=find_packages(),
       install_requires=[
           'torch',

diff --git a/torchesn/nn/echo_state_network.py b/torchesn/nn/echo_state_network.py
@@ -128,7 +128,8 @@ def forward(self, input, h_0, washout=0, target=None):
                     padded_input = padded_input.transpose(0, 1)
                 output = torch.cat([padded_input[washout:], output], -1)
             else:
-                input = input.transpose(0, 1)
+                if self.batch_first:
+                    input = input.transpose(0, 1)
                 output = torch.cat([input[washout:], output], -1)
 
         if self.readout_training == 'online' or target is None: