diff --git a/src/neural-network-gpu.ts b/src/neural-network-gpu.ts
index 57adbf9b..d7bb2b4f 100644
--- a/src/neural-network-gpu.ts
+++ b/src/neural-network-gpu.ts
@@ -36,12 +36,11 @@ function loss(
 
 function updateMemory(
   this: IKernelFunctionThis,
-  actual: number,
-  expected: number,
-  inputs: LossFunctionInputs,
   memory: NeuralNetworkMemory,
   memorySize: number,
-  loss: number
+  loss: number,
+  outputs: number[][],
+  sizes: number[]
 ) {
   const layer = this.thread.z;
   const neuron = this.thread.y;
@@ -291,10 +290,19 @@ export class NeuralNetworkGPU<
   // @ts-expect-error
   biases: KernelOutput[] = [];
 
+
+  losses: KernelOutput[] = [];
+  ram: KernelOutput = ;
+
   constructor(options: Partial<INeuralNetworkGPUOptions> = {}) {
     super(options);
     this.errorCheckInterval = 100;
     this.gpu = new GPU({ mode: options.mode });
+    this._memoryFunction = options.memory ?? DEFAULT_MEMORY_FUNCTION;
+  }
+
+  get loss(): number {
+    return this._loss;
   }
 
   get lossFunction(): LossFunction {
@@ -391,6 +399,14 @@ export class NeuralNetworkGPU<
       });
     }
 
+    // Determine whether `options.updateMemory` was passed to the constructor.
+    if (this._memoryFunction) {
+      // Get the defined memory function.
+      const updateMemory = this._memoryKernel;
+      // Update the neural network's memory matrix after each feed-forward pass.
+      updateMemory(this.outputs, this.memory, this.sizes, this.memorySize, this.loss);
+    }
+
     this.texturizeInputData = this.gpu.createKernel(
       function (value: number[]): number {
         return value[this.thread.x];
@@ -442,7 +458,6 @@ export class NeuralNetworkGPU<
     }
 
     const loss: LossFunction = this._lossFunction ?? DEFAULT_LOSS_FUNCTION;
-    const updateMemory: MemoryFunction = this._memoryFunction ?? DEFAULT_MEMORY_FUNCTION;
 
     calcDeltas = alias(
       utils.getMinifySafeName(() => calcDeltas),
@@ -523,6 +538,7 @@ export class NeuralNetworkGPU<
 
       let output;
       if (layer === this.outputLayer) {
+        // eslint-disable-next-line @typescript-eslint/ban-ts-comment
         // @ts-ignore
         output = this.backwardPropagate[layer](this.outputs[layer], target, this.outputs[0], this.memory);
       } else {
diff --git a/src/neural-network.ts b/src/neural-network.ts
index 9e3e939c..943c06da 100644
--- a/src/neural-network.ts
+++ b/src/neural-network.ts
@@ -58,7 +58,7 @@ export type NeuralNetworkActivation =
 export interface IJSONLayer {
   biases: number[];
   weights: number[][];
-  memory: number[][];
+  memory?: number[][];
 }
 
 export interface INeuralNetworkJSON {
@@ -78,6 +78,7 @@ export interface INeuralNetworkOptions {
   outputSize: number;
   binaryThresh: number;
   hiddenLayers?: number[];
+  memory?: MemoryFunction;
   memorySize: number;
 }
 
@@ -193,7 +194,7 @@ export class NeuralNetwork<
   _formatInput: NeuralNetworkFormatter | null = null;
   _formatOutput: NeuralNetworkFormatter | null = null;
 
-  _memory: NeuralNetworkMemory;
+  _memory?: NeuralNetworkMemory;
 
   runInput: (input: Float32Array) => Float32Array = (input: Float32Array) => {
     this.setActivation();
@@ -229,9 +230,23 @@ export class NeuralNetwork<
       this.sizes = [inputSize].concat(hiddenLayers ?? []).concat([outputSize]);
     }
 
+    this._memoryFunction = options.memory;
     // Initialize memory matrix
     const { memorySize } = this.options ?? 0;
-    this._memory = this.replaceMemory(memorySize);
+    if (memorySize) this._memory = this.replaceMemory(memorySize);
+  }
+
+  public get memory(): NeuralNetworkMemory | undefined {
+    return this._memory;
+  }
+
+  public get memoryFunction(): MemoryFunction | undefined {
+    return this._memoryFunction;
+  }
+
+  public get memorySize(): number {
+    if (!this._memory || !this._memory[0] || !this._memory[0][0]) return 0;
+    return this._memory[0][0].length;
   }
 
   /**
@@ -308,10 +323,6 @@ export class NeuralNetwork<
     return this.sizes.length > 0;
   }
 
-  public get memory(): NeuralNetworkMemory {
-    return this._memory;
-  }
-
   run(input: Partial<InputType>): OutputType {
     if (!this.isRunnable) {
       throw new Error('network not runnable');
@@ -1234,19 +1245,23 @@ export class NeuralNetwork<
     const jsonLayerBiases = this.biases.map((layerBiases) =>
       Array.from(layerBiases)
     );
-    const jsonLayerMemory = this.memory.map(layerMemory =>
-      layerMemory.map(
-        nodeMemory => Array.from(nodeMemory)
-      )
-    );
+    let jsonLayerMemory;
+    if (this.memory) {
+      jsonLayerMemory = this.memory.map(layerMemory =>
+        layerMemory.map(
+          nodeMemory => Array.from(nodeMemory)
+        )
+      );
+    }
     const jsonLayers: IJSONLayer[] = [];
     const outputLength = this.sizes.length - 1;
     for (let i = 0; i <= outputLength; i++) {
-      jsonLayers.push({
+      const jsonLayer: IJSONLayer = {
         weights: jsonLayerWeights[i] ?? [],
-        biases: jsonLayerBiases[i] ?? [],
-        memory: jsonLayerMemory[i] ?? []
-      });
+        biases: jsonLayerBiases[i] ?? []
+      };
+      if (jsonLayerMemory) jsonLayer.memory = jsonLayerMemory[i] ?? [];
+      jsonLayers.push(jsonLayer);
     }
     return {
       type: 'NeuralNetwork',
@@ -1290,15 +1305,18 @@ export class NeuralNetwork<
     const layerBiases = this.biases.map((layerBiases, layerIndex) =>
       Float32Array.from(jsonLayers[layerIndex].biases)
     );
-    const layerMemory = this.memory.map((memory, layerIndex) =>
-      Array.from(jsonLayers[layerIndex].memory).map(nodeMemory =>
-        Float32Array.from(nodeMemory)
-      )
-    );
+    let layerMemory;
+    if (this.memory) {
+      layerMemory = this.memory.map((memory, layerIndex) =>
+        Array.from(jsonLayers[layerIndex].memory ?? []).map(nodeMemory =>
+          Float32Array.from(nodeMemory)
+        )
+      );
+    }
     for (let i = 0; i <= this.outputLayer; i++) {
       this.weights[i] = layerWeights[i] || [];
       this.biases[i] = layerBiases[i] || [];
-      this.memory[i] = layerMemory[i] || [];
+      if (layerMemory && this.memory) this.memory[i] = (layerMemory && layerMemory[i]) ?? [];
     }
     return this;
   }
diff --git a/src/utilities/loss.ts b/src/utilities/loss.ts
index 7467745c..f861243f 100644
--- a/src/utilities/loss.ts
+++ b/src/utilities/loss.ts
@@ -14,10 +14,9 @@ export type LossFunction = (
 
 export type MemoryFunction = (
   this: IKernelFunctionThis,
-  actual: number,
-  expected: number,
-  inputs: LossFunctionInputs,
   memory: NeuralNetworkMemory,
   memorySize: number,
-  loss: number
+  loss: number,
+  outputs: number[][],
+  sizes: number[]
 ) => number;