Update LinearRegression.swift

rafaelesantos · Jun 11, 2024 · 04244d3 · 04244d3
1 parent 40ad824
commit 04244d3
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/Sources/RefdsAlgorithm/MachineLearning/LinearRegression.swift b/Sources/RefdsAlgorithm/MachineLearning/LinearRegression.swift
@@ -78,15 +78,15 @@ public class LinearRegression {
         semaphore.signal()
     }
 
-    public func predict(for features: [Double]) -> Double? {
+    public func predict(for features: [NSNumber]) -> Double? {
         semaphore.wait()
         guard !coefficients.isEmpty,
               var prediction = coefficients[safe: 0]
         else { return nil }
 
         for i in features.indices {
             if let coefficient = coefficients[safe: i + 1] {
-                prediction += coefficient * features[i]
+                prediction += coefficient * features[i].doubleValue
             }
         }
         semaphore.signal()