🐛 fix: input GPU check (#106)

owkin · Aug 9, 2023 · 424cd06 · 424cd06
1 parent bc563a7
commit 424cd06
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 2 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,6 +4,12 @@ All notable changes to this project will be documented in this file.
 
 ## [unreleased]
 
+## 0.3.1 (2023-08-09)
+
+### Bug Fixes
+
+🐛 **fix**: input GPU check ([#106](https://github.com/owkin/GrAIdient/pull/106))
+
 ## 0.3.0 (2023-08-04)
 
 ### Features

diff --git a/Sources/GrAIdient/Layer1D/Base/LayerInput1D.swift b/Sources/GrAIdient/Layer1D/Base/LayerInput1D.swift
@@ -136,7 +136,7 @@ open class LayerInput1D: Layer1D
         batchSize: Int,
         nbNeurons: Int) throws
     {
-        if data.nbElems > batchSize * nbNeurons
+        if batchSize * nbNeurons > data.nbElems
         {
             throw LayerError.DataSize
         }

diff --git a/Sources/GrAIdient/Layer2D/Base/LayerInput2D.swift b/Sources/GrAIdient/Layer2D/Base/LayerInput2D.swift
@@ -199,7 +199,7 @@ open class LayerInput2D: Layer2D
         batchSize: Int,
         nbChannels: Int, height: Int, width: Int) throws
     {
-        if data.nbElems > batchSize * nbChannels * height * width
+        if batchSize * nbChannels * height * width > data.nbElems
         {
             throw LayerError.DataSize
         }