diff --git a/roboquant/__init__.py b/roboquant/__init__.py
index 6364085..eaded7b 100644
--- a/roboquant/__init__.py
+++ b/roboquant/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "0.8.0"
+__version__ = "0.8.1"
 
 import logging
 
diff --git a/roboquant/ml/features.py b/roboquant/ml/features.py
index 29be711..a814057 100644
--- a/roboquant/ml/features.py
+++ b/roboquant/ml/features.py
@@ -49,7 +49,7 @@ def cache(self, validate=False) -> "Feature[T]":
 
     def returns(self, period=1) -> "Feature[T]":
         if period == 1:
-            return ReturnsFeature(self)  # type: ignore
+            return ReturnFeature(self)  # type: ignore
         return LongReturnsFeature(self, period)  # type: ignore
 
     def normalize(self, min_period=3) -> "Feature[T]":
@@ -359,7 +359,7 @@ def size(self) -> int:
         return len(self.assets)
 
 
-class ReturnsFeature(Feature):
+class ReturnFeature(Feature):
 
     def __init__(self, feature: Feature) -> None:
         super().__init__()
@@ -421,13 +421,12 @@ def __init__(self, feature: Feature, period: int) -> None:
     def calc(self, value):
         values = self.feature.calc(value)
         h = self.history
+        h.append(values)
 
         if len(h) < h.maxlen:  # type: ignore
-            h.append(values)
             return self._full_nan()
 
         r = max(h) / h[0] - 1.0
-        h.append(values)
         return r
 
     def size(self) -> int:
@@ -438,38 +437,6 @@ def reset(self):
         self.feature.reset()
 
 
-class MaxReturnFeature2(Feature):
-    """Calculate the maximum return over a certain period."""
-
-    def __init__(self, feature: Feature, period: int) -> None:
-        super().__init__()
-        self.feature: Feature = feature
-        self.history = np.full((period, self.size()), float("nan"), dtype=np.float32)
-        self.idx = -1
-
-    def calc(self, value):
-        self.idx += 1
-        values = self.feature.calc(value)
-        h = self.history
-        h[self.idx] = values
-
-        hist_len = len(h)
-        if self.idx < hist_len:
-            return self._full_nan()
-
-        root_idx = self.idx % hist_len + 1
-        root_idx = root_idx if root_idx < hist_len else 0
-        r = np.max(h) / h[root_idx] - 1.0
-        return r
-
-    def size(self) -> int:
-        return self.feature.size()
-
-    def reset(self):
-        self.idx = -1
-        self.feature.reset()
-
-
 class MinReturnFeature(Feature):
     """Calculate the minimum return over a certain period.
     This will only work on features that return a single value.
@@ -483,13 +450,12 @@ def __init__(self, feature: Feature, period: int) -> None:
     def calc(self, value):
         values = self.feature.calc(value)
         h = self.history
+        h.append(values)
 
         if len(h) < h.maxlen:  # type: ignore
-            h.append(values)
             return self._full_nan()
 
         r = min(h) / h[0] - 1.0
-        h.append(values)
         return r
 
     def size(self) -> int:
diff --git a/tests/unit/test_features.py b/tests/unit/test_features.py
index 0dd3cac..f676325 100644
--- a/tests/unit/test_features.py
+++ b/tests/unit/test_features.py
@@ -10,10 +10,11 @@
     NormalizeFeature,
     PriceFeature,
     SMAFeature,
-    ReturnsFeature,
+    ReturnFeature,
     VolumeFeature,
     FixedValueFeature,
     DayOfWeekFeature,
+    MaxReturnFeature
 )
 from tests.common import get_feed
 
@@ -31,11 +32,13 @@ def test_all_features(self):
             PriceFeature(symbol1, price_type="OPEN"),
             SMAFeature(FixedValueFeature(np.ones((3,))), 8),
             SMAFeature(PriceFeature(symbol2, price_type="CLOSE"), 10),
-            ReturnsFeature(PriceFeature(symbol1, price_type="OPEN")),
+            ReturnFeature(PriceFeature(symbol1, price_type="OPEN")),
+            MaxReturnFeature(PriceFeature(symbol1, price_type="CLOSE"), 4),
             VolumeFeature(symbol2),
             DayOfWeekFeature(),
         )
         channel = feed.play_background()
+        result = None
         while evt := channel.get():
             result = feature.calc(evt)
             self.assertTrue(len(result) == feature.size())