diff --git a/include/aidge/data/Tensor.hpp b/include/aidge/data/Tensor.hpp
index 5c34b0533ce63c8c14ffedfc98f458d726342946..9b2c6d76e470f55e91758b6a945d5ebcab7a9752 100644
--- a/include/aidge/data/Tensor.hpp
+++ b/include/aidge/data/Tensor.hpp
@@ -649,7 +649,7 @@ public:
     }
 
     void setGrad(std::shared_ptr<Tensor> newGrad) {
-        AIDGE_ASSERT(newGrad->dims() == mDims, "wrong dimensions");
+        AIDGE_ASSERT(!newGrad || newGrad->dims() == mDims, "wrong dimensions");
         mGrad = newGrad;
     }