diff --git a/benchmark/scripts/scanThreads/config_CPPCOFD.csv b/benchmark/scripts/scanThreads/config_CPPCOFD.csv
index 69e4c717..43fea20f 100644
--- a/benchmark/scripts/scanThreads/config_CPPCOFD.csv
+++ b/benchmark/scripts/scanThreads/config_CPPCOFD.csv
@@ -5,6 +5,6 @@ bCol,500,U64
 sketchDimension,25,U64
 ptFile,torchscripts/CRS.pt,String
 useCPP,1,U64
-cppAlgoTag,CoOFD,String
+cppAlgoTag,cooFD,String
 threads,1,U64
 forceMP,1,U64
\ No newline at end of file
diff --git a/benchmark/scripts/scanThreads/config_CPPSMPPCA.csv b/benchmark/scripts/scanThreads/config_CPPSMPPCA.csv
index ffe8d468..7c3a0223 100644
--- a/benchmark/scripts/scanThreads/config_CPPSMPPCA.csv
+++ b/benchmark/scripts/scanThreads/config_CPPSMPPCA.csv
@@ -4,6 +4,6 @@ aCol,1000,U64
 bCol,500,U64
 sketchDimension,25,U64
 threads,1,U64
-ptFile,,String
+ptFile,torchscripts/WeightedCR.pt,String
 useCPP,1,U64
 cppAlgoTag,smp-pca,String
\ No newline at end of file
diff --git a/src/CPPAlgos/SMPPCACPPAlgo.cpp b/src/CPPAlgos/SMPPCACPPAlgo.cpp
index 68d1f917..bfcf38fb 100644
--- a/src/CPPAlgos/SMPPCACPPAlgo.cpp
+++ b/src/CPPAlgos/SMPPCACPPAlgo.cpp
@@ -6,7 +6,7 @@
 
 namespace AMMBench {
 torch::Tensor AMMBench::SMPPCACPPAlgo::amm(torch::Tensor A, torch::Tensor B, uint64_t k) {
-   // Step 1: Input
+   // Step 1: Input A:n1*d B:d*n2
     A = A.t(); // d*n1
     int64_t d = A.size(0);
     int64_t n1 = A.size(1);
diff --git a/src/CPPAlgos/WeightedCRCPPAlgo.cpp b/src/CPPAlgos/WeightedCRCPPAlgo.cpp
index 59a5c75b..f6f056a1 100644
--- a/src/CPPAlgos/WeightedCRCPPAlgo.cpp
+++ b/src/CPPAlgos/WeightedCRCPPAlgo.cpp
@@ -9,13 +9,14 @@ namespace AMMBench {
 torch::Tensor AMMBench::WeightedCRCPPAlgo::amm(torch::Tensor A, torch::Tensor B, uint64_t c) {
   
   int64_t n = A.size(1); // A: m*n, B: n*d
+  std::cout << "A shape: " << A.sizes() << std::endl;
+  std::cout << "B shape: " << B.sizes() << std::endl;
 
   // Probability distribution
-  torch::Tensor probability_distribution = torch::zeros({n});
-    for (int i = 0; i < n; ++i) {
-      probability_distribution[i] = torch::norm(A.t()[i]) * torch::norm(B[i]);
-    }
-    probability_distribution /= probability_distribution.sum();
+  torch::Tensor col_norm_A = torch::norm(A, /*p=*/2, /*dim=*/0); // norm on columns of A
+  torch::Tensor row_norm_B = torch::norm(B, /*p=*/2, /*dim=*/1); // norm on rows of B
+  torch::Tensor probability_distribution = torch::mul(col_norm_A, row_norm_B);
+  probability_distribution /= probability_distribution.sum();
 
   // S
   torch::Tensor sample_indices = torch::multinomial(probability_distribution, /*num_samples*/c, /*replacement*/true);
@@ -24,11 +25,8 @@ torch::Tensor AMMBench::WeightedCRCPPAlgo::amm(torch::Tensor A, torch::Tensor B,
   std::tie(unique_indices, _, occurences) = at::_unique2(sample_indices, /*sorted*/false, /*return_inverse*/false, /*return_counts*/true);
 
   torch::Tensor S = torch::zeros({n, unique_indices.size(0)});
-
-  for (int trial = 0; trial < unique_indices.size(0); ++trial) {
-    int index = unique_indices[trial].item<int>();
-    S[index][trial] = 1;
-  }
+  torch::Tensor trial = torch::arange(unique_indices.size(0));
+  S.index_put_({unique_indices, trial}, 1);
 
   // D
   torch::Tensor D = torch::diag(torch::sqrt(occurences) / torch::sqrt(torch::tensor({static_cast<long>(c)}, torch::kLong) * probability_distribution.index_select(0, unique_indices)));
diff --git a/test/SystemTest/WeightedCRTest.cpp b/test/SystemTest/WeightedCRTest.cpp
index 1466ce5b..9c4b915b 100644
--- a/test/SystemTest/WeightedCRTest.cpp
+++ b/test/SystemTest/WeightedCRTest.cpp
@@ -58,10 +58,11 @@ TEST_CASE("Test Weighted CR in cpp", "[short]")
 {
     torch::manual_seed(114514);
     AMMBench::WeightedCRCPPAlgo wcr;
-    auto A = torch::rand({400, 400});
-    auto B = torch::rand({400, 400});
+    auto A = torch::rand({500, 400});
+    auto B = torch::rand({400, 600});
     auto realC = torch::matmul(A, B);
     auto ammC = wcr.amm(A, B, 20);
     double froError = INTELLI::UtilityFunctions::relativeFrobeniusNorm(realC, ammC);
+    std::cout << "froError: " << froError << std::endl;
     REQUIRE(froError < 0.5);
 }
\ No newline at end of file