diff --git a/CHANGELOG.md b/CHANGELOG.md
index c6f5db04..4564bb16 100644
--- 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
index d34977af..fbc22d41 100644
--- 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
index ed3532f7..3d8caf2c 100644
--- 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
         }