Merge pull request #14 from dahong67/restrict-lossfunctions

Restrict LossFunctions.jl support to DistanceLoss and MarginLoss
dahong67 · Oct 10, 2023 · b0b1a3b · b0b1a3b
2 parents c4fa875 + 2db3cd3
commit b0b1a3b
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/ext/LossFunctionsExt.jl b/ext/LossFunctionsExt.jl
@@ -2,6 +2,8 @@ module LossFunctionsExt
 
 using GCPDecompositions, LossFunctions
 
+const SupportedLosses = Union{LossFunctions.DistanceLoss,LossFunctions.MarginLoss}
+
 """
     gcp(X::Array, r, loss::LossFunctions.SupervisedLoss[, lower]) -> CPD
 
@@ -15,7 +17,7 @@ with respect to the loss function `loss` and return a `CPD` object.
   - `loss` : loss function from LossFunctions.jl
   - `lower` : lower bound for factor matrix entries, `default = -Inf`
 """
-GCPDecompositions.gcp(X::Array, r, loss::LossFunctions.SupervisedLoss, lower = -Inf) =
+GCPDecompositions.gcp(X::Array, r, loss::SupportedLosses, lower = -Inf) =
     GCPDecompositions._gcp(
         X,
         r,