diff --git a/src/TR.jl b/src/TR.jl
index 5b906b1..77fe209 100644
--- a/src/TR.jl
+++ b/src/TR.jl
@@ -62,7 +62,7 @@ elseif regType == "legendre"
     P, _    = plegendre(regParam1-1, p);   
     regMat[end-regParam1+1:end,:] = sqrt(regParam2) * P;
 elseif regType == "L2"
-    regMat = I(size(X,2));
+    regMat = I(p);
 elseif regType == "std"
     regMat = regParam2;
 end