diff --git a/lib/ruby_linear_regression.rb b/lib/ruby_linear_regression.rb
index 6842189..2022e08 100644
--- a/lib/ruby_linear_regression.rb
+++ b/lib/ruby_linear_regression.rb
@@ -28,7 +28,7 @@ def load_training_data x_data, y_data, normalize = true
         @x = Matrix.rows( x_data )
         @y = Matrix.rows( y_data.collect { |e| [e] } )
 
-        @theta = Matrix.zero(@x.column_count, 1)
+        @theta = Matrix.zero(@x.column_size, 1)
   end
 
   # Compute the mean squared cost / error function