chore: fix test

mzbac · Oct 13, 2024 · 1a3d163 · 1a3d163
1 parent 3021ed0
commit 1a3d163
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/Package.swift b/Package.swift
@@ -11,7 +11,7 @@ let package = Package(
       targets: ["FluxSwift"])
   ],
   dependencies: [
-    .package(url: "https://github.com/ml-explore/mlx-swift", from: "0.16.0"),
+    .package(url: "https://github.com/ml-explore/mlx-swift", from: "0.18.0"),
     .package(url: "https://github.com/huggingface/swift-transformers", from: "0.1.13"),
   ],
   targets: [

diff --git a/Sources/FluxConfiguration.swift b/Sources/FluxConfiguration.swift
@@ -103,7 +103,7 @@ func fuseLoraWeights(
         let transformerWeight = fusedWeights[weightKey]
       {
         let loraScale: Float = 1.0 
-        let loraFused = MLX.matmul(loraB, loraA)
+        let loraFused = matmul(loraB, loraA)
         fusedWeights[weightKey] = transformerWeight + loraScale * loraFused
       }
     }

diff --git a/Tests/MultiModalDiffusionTransformerTests.swift b/Tests/MultiModalDiffusionTransformerTests.swift
@@ -16,7 +16,7 @@ final class MultiModalDiffusionTransformerTests: XCTestCase {
         let pooledPromptEmbeds = MLXRandom.normal([1, config.pooledProjectionDim])
         let hiddenStates = MLXRandom.normal([1, config.jointAttentionDim, config.inChannels])
 
-        let evaluateParameters = EvaluateParameters()
+        let evaluateParameters = EvaluateParameters(width: 1024, height: 1024, shiftSigmas: true)
 
         let output = transformer(
             t: t,