diff --git a/src/stream/stream/partition.rs b/src/stream/stream/partition.rs
index 077d4ec2..74231ab7 100644
--- a/src/stream/stream/partition.rs
+++ b/src/stream/stream/partition.rs
@@ -44,15 +44,12 @@ where
 
             match next {
                 Some(v) => {
-                    let mut res = this.res.take().unwrap();
-
+                    let res = this.res.as_mut().unwrap();
                     if (this.f)(&v) {
                         res.0.extend(Some(v))
                     } else {
                         res.1.extend(Some(v))
                     }
-
-                    *this.res = Some(res);
                 }
                 None => return Poll::Ready(this.res.take().unwrap()),
             }