Reflected same change through to python BCF

Author: andrewherren

stochtree/bcf.py

@@ -1093,14 +1093,6 @@ def sample(
         else:
             variable_subset_variance = [i for i in range(X_train.shape[1])]
 
-        # Set num_features_subsample to default, ncol(X_train), if not already set
-        if num_features_subsample_mu is None:
-            num_features_subsample_mu = X_train.shape[1]
-        if num_features_subsample_tau is None:
-            num_features_subsample_tau = X_train.shape[1]
-        if num_features_subsample_variance is None:
-            num_features_subsample_variance = X_train.shape[1]
-
         # Determine whether a test set is provided
         self.has_test = X_test is not None
 
@@ -1498,6 +1490,14 @@ def sample(
         # Store propensity score requirements of the BCF forests
         self.propensity_covariate = propensity_covariate
 
+        # Set num_features_subsample to default, ncol(X_train), if not already set
+        if num_features_subsample_mu is None:
+            num_features_subsample_mu = X_train_processed.shape[1]
+        if num_features_subsample_tau is None:
+            num_features_subsample_tau = X_train_processed.shape[1]
+        if num_features_subsample_variance is None:
+            num_features_subsample_variance = X_train_processed.shape[1]
+
         # Container of variance parameter samples
         self.num_gfr = num_gfr
         self.num_burnin = num_burnin