diff --git a/src/main/scala/leon/utils/UnitElimination.scala b/src/main/scala/leon/utils/UnitElimination.scala
index e5672bd5128f8fabc7f64b4f04038ff78b39f307..60f8bc232047c47900f728d817ac285d06f8482b 100644
--- a/src/main/scala/leon/utils/UnitElimination.scala
+++ b/src/main/scala/leon/utils/UnitElimination.scala
@@ -97,10 +97,8 @@ object UnitElimination extends TransformationPhase {
           val (newFd, rest) = if(fd.params.exists(vd => vd.getType == UnitType)) {
             val freshFunDef = new FunDef(FreshIdentifier(fd.id.name), fd.tparams, fd.returnType, fd.params.filterNot(vd => vd.getType == UnitType), fd.defType).setPos(fd)
             freshFunDef.addAnnotation(fd.annotations.toSeq:_*)
-            freshFunDef.precondition = fd.precondition //TODO: maybe removing unit from the conditions as well..
-            freshFunDef.postcondition = fd.postcondition//TODO: maybe removing unit from the conditions as well..
             fun2FreshFun += (fd -> freshFunDef)
-            freshFunDef.body = fd.body.map(b => removeUnit(b))
+            freshFunDef.fullBody = removeUnit(fd.fullBody)
             val restRec = removeUnit(b)
             fun2FreshFun -= fd
             (freshFunDef, restRec)