diff --git a/get_deps.sh b/get_deps.sh
index 4a70f536e..a5d0f07e1 100755
--- a/get_deps.sh
+++ b/get_deps.sh
@@ -12,7 +12,7 @@ HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)"
 if [[ $1 == --help || $1 == help ]]; then
 	cat <<-END
 		[ARGVARS...] get_deps.sh [cpu|gpu] [--help|help]
-		
+
 		Argument variables:
 		CPU=1              Get CPU dependencies
 		GPU=1              Get GPU dependencies
@@ -103,7 +103,7 @@ if [[ $WITH_TF != 0 ]]; then
 			fi
 			if [[ $ARCH == x64 ]]; then
 				TF_ARCH=x86_64
-				
+
 				LIBTF_URL_BASE=https://storage.googleapis.com/tensorflow/libtensorflow
 			elif [[ $ARCH == arm64v8 ]]; then
 				TF_ARCH=arm64
@@ -134,7 +134,7 @@ if [[ $WITH_TF != 0 ]]; then
 		mkdir $LIBTENSORFLOW.x
 		tar xf $LIBTF_ARCHIVE --no-same-owner -C $LIBTENSORFLOW.x
 		mv $LIBTENSORFLOW.x $LIBTENSORFLOW
-		
+
 		echo "Done."
 	else
 		echo "TensorFlow is in place."
@@ -152,7 +152,7 @@ if [[ $WITH_TFLITE != 0 ]]; then
 
 	if [[ ! -d $LIBTFLITE ]]; then
 		echo "Installing TensorFlow Lite ..."
-		
+
 		LIBTF_URL_BASE=https://s3.amazonaws.com/redismodules/tensorflow
 		if [[ $OS == linux ]]; then
 			TFLITE_OS="linux"
@@ -183,7 +183,7 @@ if [[ $WITH_TFLITE != 0 ]]; then
 		mkdir $LIBTFLITE.x
 		tar xf $LIBTFLITE_ARCHIVE --no-same-owner -C $LIBTFLITE.x
 		mv $LIBTFLITE.x $LIBTFLITE
-		
+
 		echo "Done."
 	else
 		echo "TensorFlow Lite is in place."
@@ -241,14 +241,14 @@ if [[ $WITH_PT != 0 ]]; then
 
 			[[ ! -f $LIBTORCH_ARCHIVE || $FORCE == 1 ]] && wget -q $LIBTORCH_URL
 		fi
-		
+
 		rm -rf $LIBTORCH.x
 		mkdir $LIBTORCH.x
 
 		tar xf $LIBTORCH_ARCHIVE --no-same-owner -C $LIBTORCH.x
 		mv $LIBTORCH.x/libtorch $LIBTORCH
 		rmdir $LIBTORCH.x
-		
+
 		echo "Done."
 	else
 		echo "libtorch is in place."
@@ -282,7 +282,7 @@ fi
 
 ################################################################################### ONNXRUNTIME
 
-ORT_VERSION="1.6.0"
+ORT_VERSION="1.7.1"
 
 if [[ $WITH_ORT != 0 ]]; then
 	[[ $FORCE == 1 ]] && rm -rf $ONNXRUNTIME
@@ -297,15 +297,13 @@ if [[ $WITH_ORT != 0 ]]; then
 			else
 				ORT_BUILD="-gpu"
 			fi
+			ORT_URL_BASE=https://s3.amazonaws.com/redismodules/onnxruntime
 			if [[ $ARCH == x64 ]]; then
 				ORT_ARCH=x64
-				ORT_URL_BASE=https://github.com/microsoft/onnxruntime/releases/download/v${ORT_VERSION}
 			elif [[ $ARCH == arm64v8 ]]; then
 				ORT_ARCH=arm64
-				ORT_URL_BASE=https://s3.amazonaws.com/redismodules/onnxruntime
 			elif [[ $ARCH == arm32v7 ]]; then
 				ORT_ARCH=arm
-				ORT_URL_BASE=https://s3.amazonaws.com/redismodules/onnxruntime
 			fi
 		elif [[ $OS == macos ]]; then
 			ORT_OS=osx
@@ -322,7 +320,7 @@ if [[ $WITH_ORT != 0 ]]; then
 		mkdir $ONNXRUNTIME.x
 		tar xzf ${ORT_ARCHIVE} --no-same-owner --strip-components=1 -C $ONNXRUNTIME.x
 		mv $ONNXRUNTIME.x $ONNXRUNTIME
-		
+
 		echo "Done."
 	else
 		echo "ONNXRuntime is in place."
diff --git a/opt/build/onnxruntime/Dockerfile.arm7 b/opt/build/onnxruntime/Dockerfile.arm7
index 732a9c127..4cfea0af0 100755
--- a/opt/build/onnxruntime/Dockerfile.arm7
+++ b/opt/build/onnxruntime/Dockerfile.arm7
@@ -4,6 +4,7 @@ ARG ONNXRUNTIME_REPO=https://github.com/Microsoft/onnxruntime
 ARG ONNXRUNTIME_BRANCH=rel-1.0.0
 ARG ONNXRUNTIME_VER=1.0.0
 ARG ARCH_FLAG="--arm"
+ARG ARCH=arm64v7
 
 RUN [ "cross-build-start" ]
 
diff --git a/opt/build/onnxruntime/Dockerfile.x64 b/opt/build/onnxruntime/Dockerfile.x64
index a16576fb2..439899269 100755
--- a/opt/build/onnxruntime/Dockerfile.x64
+++ b/opt/build/onnxruntime/Dockerfile.x64
@@ -5,30 +5,36 @@ ARG OS=debian:buster
 #----------------------------------------------------------------------------------------------
 FROM ${OS}
 
-ARG ONNXRUNTIME_REPO=https://github.com/Microsoft/onnxruntime
-ARG ONNXRUNTIME_VER=1.0.0
+ARG ONNXRUNTIME_REPO=https://github.com/RedisAI/onnxruntime
+ARG ONNXRUNTIME_VER=1.7.1
+ARG ARCH=x64
 
 RUN apt-get -qq update
-RUN apt-get -qq install -y curl wget tar git
-RUN apt-get -qq install -y build-essential cmake
-RUN apt-get -qq install -y libcurl4-openssl-dev libssl-dev libatlas-base-dev zlib1g-dev
-
-RUN apt-get -qq install -y python3 python3-pip python3-dev
+RUN apt-get -qq install -y curl wget tar git \
+            build-essential cmake \
+            libcurl4-openssl-dev libssl-dev libatlas-base-dev zlib1g-dev \
+            python3 python3-pip python3-dev python3-numpy
 RUN pip3 install --upgrade pip setuptools wheel
-# RUN pip3 install numpy
-RUN apt-get -q install -y python3-numpy
+
+ENV LANG=en_US.UTF-8
+RUN apt-get install -y locales && \
+    sed -i -e "s/# $LANG.*/$LANG UTF-8/" /etc/locale.gen && \
+    dpkg-reconfigure --frontend=noninteractive locales && \
+    update-locale LANG=$LANG
 
 WORKDIR /build
 
-ADD ./pack.sh /build/
 ARG BUILDTYPE=MinSizeRel
+
 ARG BUILDARGS="--config ${BUILDTYPE} --parallel"
 
-RUN git clone --single-branch --branch rel-${ONNXRUNTIME_VER} --recursive ${ONNXRUNTIME_REPO} onnxruntime
-	
-RUN	cd onnxruntime ;\
-	./build.sh ${BUILDARGS} --update --build ;\
-	./build.sh ${BUILDARGS} --build_shared_lib
+RUN git clone --single-branch --branch rel-${ONNXRUNTIME_VER} ${ONNXRUNTIME_REPO} onnxruntime
+WORKDIR /build/onnxruntime
+RUN git fetch --recurse-submodules -j4
+RUN ./build.sh ${BUILDARGS} --update --build
+RUN ./build.sh ${BUILDARGS} --build_shared_lib
 # RUN ./build.sh ${BUILDARGS} --enable_pybind --build_wheel
 
-RUN ./pack.sh ${ONNXRUNTIME_VER}
+ADD ./pack.sh /build
+WORKDIR /build
+RUN ./pack.sh ${ONNXRUNTIME_VER} ${ARCH}
diff --git a/opt/build/onnxruntime/Dockerfile.x64-gpu b/opt/build/onnxruntime/Dockerfile.x64-gpu
new file mode 100644
index 000000000..e691f7f4a
--- /dev/null
+++ b/opt/build/onnxruntime/Dockerfile.x64-gpu
@@ -0,0 +1,43 @@
+ARG OS=ubuntu18.04
+ARG CUDA_VER=11.0-cudnn8
+
+FROM nvidia/cuda:11.0-cudnn8-devel-ubuntu18.04
+
+ARG ONNXRUNTIME_REPO=https://github.com/RedisAI/onnxruntime
+ARG ONNXRUNTIME_VER=1.7.1
+ARG ARCH=x64-gpu
+
+RUN apt-get -qq update
+RUN apt-get -qq install -y curl \
+            wget tar git build-essential \
+            libcurl4-openssl-dev libssl-dev \
+            libatlas-base-dev zlib1g-dev python3 python3-pip \
+            python3-dev python3-numpy rsync
+
+RUN pip3 install --upgrade pip setuptools wheel
+
+RUN wget -q https://github.com/Kitware/CMake/releases/download/v3.19.5/cmake-3.19.5-Linux-x86_64.tar.gz -O /tmp/cmake.tgz
+WORKDIR /tmp
+RUN tar -xpf cmake.tgz
+RUN rsync -aqH cmake*/* /usr
+
+
+ENV LANG=en_US.UTF-8
+RUN apt-get install -y locales && \
+    sed -i -e "s/# $LANG.*/$LANG UTF-8/" /etc/locale.gen && \
+    dpkg-reconfigure --frontend=noninteractive locales && \
+    update-locale LANG=$LANG
+
+WORKDIR /build
+
+ARG BUILDTYPE=MinSizeRel
+ARG BUILDARGS="--config ${BUILDTYPE} --parallel"
+
+RUN git clone --single-branch --branch rel-${ONNXRUNTIME_VER} ${ONNXRUNTIME_REPO} onnxruntime
+WORKDIR /build/onnxruntime
+RUN git fetch --recurse-submodules -j4
+RUN ./build.sh ${BUILDARGS} --update --build --use_cuda --cudnn_home /usr/local/cuda --cuda_home /usr/local/cuda --build_shared_lib --parallel
+
+ADD ./pack.sh /build/
+WORKDIR /build
+RUN ./pack.sh ${ONNXRUNTIME_VER} ${ARCH}
diff --git a/opt/build/onnxruntime/Makefile b/opt/build/onnxruntime/Makefile
index 6e785c455..1faaabae9 100755
--- a/opt/build/onnxruntime/Makefile
+++ b/opt/build/onnxruntime/Makefile
@@ -1,7 +1,7 @@
 
 ROOT=.
 
-VERSION ?= 1.0.0
+VERSION ?= 1.7.1
 OSNICK ?= buster
 
 #----------------------------------------------------------------------------------------------
@@ -16,7 +16,7 @@ DOCKER_OS.bionic=ubuntu:bionic
 DOCKER_OS.stretch=debian:stretch-slim
 DOCKER_OS.buster=debian:buster-slim
 DOCKER_OS=$(DOCKER_OS.$(OSNICK))
- 
+
 #----------------------------------------------------------------------------------------------
 
 define targets # (1=OP, 2=op)
@@ -40,10 +40,23 @@ IID_$(1)=$(1)_$(VERSION).iid
 CID_$(1)=$(1)_$(VERSION).cid
 
 build_x64:
-	@docker build --iidfile $$(IID_$(1)) -t redisfab/$(STEM)-$(1):$(VERSION) -f Dockerfile.x64 \
+	@docker build --iidfile $$(IID_$(1)) -t redisfab/$(STEM)-x64:$(VERSION) -f Dockerfile.x64 \
 		--build-arg OS=$(DOCKER_OS) $(ROOT)
 	@docker create --cidfile $$(CID_$(1)) `cat $$(IID_$(1))`
-	@docker cp `cat $$(CID_$(1))`:/build/$(STEM)-$(2)-$(VERSION).tgz .
+	@docker cp `cat $$(CID_$(1))`:/build/$(STEM)-x64-$(VERSION).tgz .
+
+.PHONY: build_x64
+endef
+
+define build_x64-gpu # (1=arch, 2=tar-arch)
+IID_$(1)=$(1)_$(VERSION).iid
+CID_$(1)=$(1)_$(VERSION).cid
+
+build_x64-gpu:
+	@docker build --iidfile $$(IID_$(1)) -t redisfab/$(STEM)-$(1):$(VERSION) -f Dockerfile.x64-gpu \
+		--build-arg OS=$(DOCKER_OS) $(ROOT)
+	@docker create --cidfile $$(CID_$(1)) `cat $$(IID_$(1))`
+	@docker cp `cat $$(CID_$(1))`:/build/$(STEM)-$(1)-$(VERSION).tgz .
 
 .PHONY: build_x64
 endef
@@ -68,11 +81,18 @@ endef
 
 define publish_x64 # (1=arch, 2=tar-arch)
 publish_x64:
-	@aws s3 cp $(STEM)-$(2)-$(VERSION).tar.gz s3://$(S3_URL)/ --acl public-read
+	@aws s3 cp $(STEM)-$(2)-$(VERSION).tgz s3://$(S3_URL)/ --acl public-read
 
 .PHONY: publish_x64
 endef
 
+define publish_x64-gpu # (1=arch, 2=tar-arch)
+publish_x64-gpu:
+	@aws s3 cp $(STEM)-$(2)-$(VERSION).tgz s3://$(S3_URL)/ --acl public-read
+
+.PHONY: publish_x64-gpu
+endef
+
 define publish_arm # (1=arch, 2=tar-arch)
 publish_$(1):
 	@aws s3 cp $(STEM)-$(2)-$(VERSION).tgz s3://$(S3_URL)/ --acl public-read
@@ -90,6 +110,7 @@ all: build publish
 build: $(BUILD_TARGETS)
 
 $(eval $(call build_x64,x64,x86_64))
+$(eval $(call build_x64-gpu,x64-gpu,x86_64-gpu))
 $(eval $(call build_arm,arm64v8,arm64))
 $(eval $(call build_arm,arm32v7,arm))
 
@@ -101,7 +122,8 @@ endif
 
 publish: $(PUBLISH_TARGETS)
 
-$(eval $(call publish_x64,x64,x86_64))
+$(eval $(call publish_x64,x64,x64))
+$(eval $(call publish_x64-gpu,x64-gpu,x64-gpu))
 $(eval $(call publish_arm,arm64v8,arm64))
 $(eval $(call publish_arm,arm32v7,arm))
 
diff --git a/opt/build/onnxruntime/pack.sh b/opt/build/onnxruntime/pack.sh
index 52eb55cb7..01b6bb5ed 100755
--- a/opt/build/onnxruntime/pack.sh
+++ b/opt/build/onnxruntime/pack.sh
@@ -2,10 +2,11 @@
 
 set -e
 VER="$1"
+PLATFORM="$2"
 
 mkdir -p pack/include pack/lib
 cp onnxruntime/build/Linux/MinSizeRel/libonnxruntime.so.${VER} pack/lib/
-cp onnxruntime/docs/C_API.md pack/
+cp onnxruntime/docs/C_API_Guidelines.md pack/
 cp onnxruntime/LICENSE pack/
 cp onnxruntime/README.md pack/
 cp onnxruntime/ThirdPartyNotices.txt pack/
@@ -14,8 +15,9 @@ cd onnxruntime/
 git rev-parse HEAD > ../pack/GIT_COMMIT_ID
 cd ..
 cp onnxruntime/include/onnxruntime/core/session/onnxruntime_c_api.h pack/include/
+cp onnxruntime/include/onnxruntime/core/providers/cuda/cuda_provider_factory.h pack/include/
 cd pack/lib/
 ln -s libonnxruntime.so.${VER} libonnxruntime.so
 cd ../..
-mv pack onnxruntime-linux-arm64-${VER}
-tar czf onnxruntime-linux-arm64-${VER}.tgz onnxruntime-linux-arm64-${VER}/
+mv pack onnxruntime-linux-${PLATFORM}-${VER}
+tar czf onnxruntime-linux-${PLATFORM}-${VER}.tgz onnxruntime-linux-${PLATFORM}-${VER}/
diff --git a/opt/redis_valgrind.sup b/opt/redis_valgrind.sup
index 5431acaf1..8b00a448c 100644
--- a/opt/redis_valgrind.sup
+++ b/opt/redis_valgrind.sup
@@ -12,13 +12,6 @@
    obj:*/libtensorflow_framework.so.*
 }
 
-{
-   ignore_unversioned_libs
-   Memcheck:Leak
-   ...
-   obj:*/libonnxruntime.so.*
-}
-
 {
    ignore_unversioned_libs
    Memcheck:Leak
diff --git a/src/backends.c b/src/backends.c
index e9e3e4d7c..f647fa585 100644
--- a/src/backends.c
+++ b/src/backends.c
@@ -451,10 +451,27 @@ int RAI_LoadBackend_ONNXRuntime(RedisModuleCtx *ctx, const char *path) {
         return REDISMODULE_ERR;
     }
 
-    RAI_backends.onnx = backend;
+    backend.get_memory_info =
+        (unsigned long long (*)(void))(unsigned long)dlsym(handle, "RAI_GetMemoryInfoORT");
+    if (backend.get_memory_info == NULL) {
+        dlclose(handle);
+        RedisModule_Log(ctx, "warning",
+                        "Backend does not export RAI_GetMemoryInfoORT. ONNX backend "
+                        "not loaded from %s",
+                        path);
+    }
+    backend.get_memory_access_num =
+        (unsigned long long (*)(void))(unsigned long)dlsym(handle, "RAI_GetMemoryAccessORT");
+    if (backend.get_memory_access_num == NULL) {
+        dlclose(handle);
+        RedisModule_Log(ctx, "warning",
+                        "Backend does not export RAI_GetMemoryAccessORT. ONNX backend "
+                        "not loaded from %s",
+                        path);
+    }
 
+    RAI_backends.onnx = backend;
     RedisModule_Log(ctx, "notice", "ONNX backend loaded from %s", path);
-
     return REDISMODULE_OK;
 }
 
diff --git a/src/backends.h b/src/backends.h
index 9c2b79231..045526955 100644
--- a/src/backends.h
+++ b/src/backends.h
@@ -77,6 +77,11 @@ typedef struct RAI_LoadedBackend {
     // Returns the backend version.
     const char *(*get_version)(void);
 
+    // Returns the backend's memory usage for INFO report
+    unsigned long long (*get_memory_info)(void);
+
+    // Returns the number of times that Redis accessed backend allocator.
+    unsigned long long (*get_memory_access_num)(void);
 } RAI_LoadedBackend;
 
 typedef struct RAI_LoadedBackends {
diff --git a/src/backends/onnxruntime.c b/src/backends/onnxruntime.c
index 51a13ca12..21494d9b0 100644
--- a/src/backends/onnxruntime.c
+++ b/src/backends/onnxruntime.c
@@ -4,19 +4,122 @@
 #include "backends/util.h"
 #include "tensor.h"
 #include "util/arr_rm_alloc.h"
+#include <stdatomic.h>
 
 #include "onnxruntime_c_api.h"
 
+// Use as a wrapper for ORT api call. If ORT api hasn't returned null, it has failed.
+// A label "error" must exist in every function that uses this macro.
+#define ONNX_VALIDATE_STATUS(x)                                                                    \
+    if ((status = (x)) != NULL)                                                                    \
+        goto error;
+
+OrtEnv *env = NULL;
+
+// For model that run on GPU, onnx will not use the custom allocator (redis allocator), but
+// the onnx allocator for GPU. But for the auxiliary allocations of the input and output names,
+// we will use the custom global allocator for models that run on GPU as well.
+OrtMemoryInfo *mem_info = NULL;
+OrtAllocator *global_allocator = NULL;
+unsigned long long OnnxMemory = 0;
+unsigned long long OnnxMemoryAccessCounter = 0;
+
+const OrtMemoryInfo *AllocatorInfo(const OrtAllocator *allocator) {
+    (void)allocator;
+    const OrtApi *ort = OrtGetApiBase()->GetApi(1);
+    if (mem_info != NULL) {
+        return mem_info;
+    }
+    if (ort->CreateCpuMemoryInfo(OrtDeviceAllocator, OrtMemTypeDefault, &mem_info) != NULL) {
+        return NULL;
+    }
+    return mem_info;
+}
+
+// Allocate address with 64-byte alignment to cope with onnx optimizations.
+void *AllocatorAlloc(OrtAllocator *ptr, size_t size) {
+
+    (void)ptr;
+    // Allocate an additional 63 bytes to ensure that we can return an address which is
+    // 64-byte aligned, and an additional space in the size of a pointer to store
+    // the address that RedisModule_Alloc returns.
+    int offset = 63 + sizeof(void *);
+    void *allocated_address = (void *)RedisModule_Alloc(size + offset);
+    size_t allocated_size = RedisModule_MallocSize(allocated_address);
+    // Update the total number of bytes that onnx is using and the number of accesses
+    // that onnx made to the allocator.
+    atomic_fetch_add(&OnnxMemory, allocated_size);
+    atomic_fetch_add(&OnnxMemoryAccessCounter, 1);
+    // This operation guarantees that p2 is the closest 64-aligned address to (p1+size_t).
+    void **aligned_address = (void **)(((size_t)(allocated_address) + offset) & (~63));
+    // This stores the address p1 right before p2 (so we can retrieve it when we free).
+    aligned_address[-1] = allocated_address;
+    return aligned_address;
+}
+
+void AllocatorFree(OrtAllocator *ptr, void *aligned_address) {
+    (void)ptr;
+    if (aligned_address == NULL) {
+        return;
+    }
+    // Retrieve the address that we originally received from RedisModule_Alloc
+    // (this is the address that we need to sent to RedisModule_Free).
+    void *allocated_address = ((void **)aligned_address)[-1];
+    size_t allocated_size = RedisModule_MallocSize(allocated_address);
+    // Update the total number of bytes that onnx is using and the number of accesses
+    // that onnx made to the allocator.
+    atomic_fetch_sub(&OnnxMemory, allocated_size);
+    atomic_fetch_add(&OnnxMemoryAccessCounter, 1);
+    return RedisModule_Free(allocated_address);
+}
+
+unsigned long long RAI_GetMemoryInfoORT() { return OnnxMemory; }
+
+unsigned long long RAI_GetMemoryAccessORT() { return OnnxMemoryAccessCounter; }
+
 int RAI_InitBackendORT(int (*get_api_fn)(const char *, void *)) {
     get_api_fn("RedisModule_Alloc", ((void **)&RedisModule_Alloc));
     get_api_fn("RedisModule_Calloc", ((void **)&RedisModule_Calloc));
     get_api_fn("RedisModule_Free", ((void **)&RedisModule_Free));
     get_api_fn("RedisModule_Realloc", ((void **)&RedisModule_Realloc));
     get_api_fn("RedisModule_Strdup", ((void **)&RedisModule_Strdup));
-
+    get_api_fn("RedisModule_Log", ((void **)&RedisModule_Log));
+    get_api_fn("RedisModule_GetThreadSafeContext", ((void **)&RedisModule_GetThreadSafeContext));
+    get_api_fn("RedisModule_FreeThreadSafeContext", ((void **)&RedisModule_FreeThreadSafeContext));
+    get_api_fn("RedisModule_MallocSize", ((void **)&RedisModule_MallocSize));
     return REDISMODULE_OK;
 }
 
+bool setDeviceId(const char *devicestr, OrtSessionOptions *session_options, RAI_Error *error) {
+
+    RAI_Device device;
+    int64_t deviceid;
+    if (!parseDeviceStr(devicestr, &device, &deviceid)) {
+        RAI_SetError(error, RAI_EMODELCREATE, "ERR unsupported device");
+        return false;
+    }
+
+    // TODO: we will need to propose a more dynamic way to request a specific provider,
+    // e.g. given the name, in ONNXRuntime
+#if RAI_ONNXRUNTIME_USE_CUDA
+    if (device == RAI_DEVICE_GPU) {
+        if (deviceid == -1) {
+            // ORT does not like device id as -1
+            deviceid = 0;
+        }
+        OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, deviceid);
+    }
+#else
+    // TODO: Do dynamic device/provider check with GetExecutionProviderType or something on these
+    // lines
+    if (device == RAI_DEVICE_GPU) {
+        RAI_SetError(error, RAI_EMODELCREATE, "ERR GPU requested but ONNX couldn't find CUDA");
+        return false;
+    }
+#endif
+    return true;
+}
+
 ONNXTensorElementDataType RAI_GetOrtDataTypeFromDL(DLDataType dtype) {
     if (dtype.code == kDLFloat) {
         switch (dtype.bits) {
@@ -74,23 +177,12 @@ DLDataType RAI_GetDLDataTypeFromORT(ONNXTensorElementDataType dtype) {
     default:
         return (DLDataType){.bits = 0};
     }
-    return (DLDataType){.bits = 0};
 }
 
 OrtValue *RAI_OrtValueFromTensors(RAI_Tensor **ts, size_t count, RAI_Error *error) {
     OrtStatus *status = NULL;
     const OrtApi *ort = OrtGetApiBase()->GetApi(1);
 
-    if (count == 0) {
-        return NULL;
-    }
-
-    OrtAllocator *allocator;
-    status = ort->GetAllocatorWithDefaultOptions(&allocator);
-    if (status != NULL) {
-        return NULL;
-    }
-
     size_t batch_size = 0;
     size_t batch_byte_size = 0;
 
@@ -100,52 +192,36 @@ OrtValue *RAI_OrtValueFromTensors(RAI_Tensor **ts, size_t count, RAI_Error *erro
     }
 
     RAI_Tensor *t0 = ts[0];
-
     const int ndim = t0->tensor.dl_tensor.ndim;
     int64_t batched_shape[ndim];
-
     for (size_t i = 1; i < ndim; i++) {
         batched_shape[i] = t0->tensor.dl_tensor.shape[i];
     }
-
     batched_shape[0] = batch_size;
 
     OrtValue *out;
-
     if (count > 1) {
-        status =
-            ort->CreateTensorAsOrtValue(allocator, batched_shape, t0->tensor.dl_tensor.ndim,
-                                        RAI_GetOrtDataTypeFromDL(t0->tensor.dl_tensor.dtype), &out);
-        if (status != NULL) {
-            goto error;
-        }
+        ONNX_VALIDATE_STATUS(
+            ort->CreateTensorAsOrtValue(global_allocator, batched_shape, t0->tensor.dl_tensor.ndim,
+                                        RAI_GetOrtDataTypeFromDL(t0->tensor.dl_tensor.dtype), &out))
 
         char *ort_data;
-        status = ort->GetTensorMutableData(out, (void **)&ort_data);
-        if (status != NULL) {
-            goto error;
-        }
-
+        ONNX_VALIDATE_STATUS(ort->GetTensorMutableData(out, (void **)&ort_data))
         size_t offset = 0;
         for (size_t i = 0; i < count; i++) {
             memcpy(ort_data + offset, RAI_TensorData(ts[i]), RAI_TensorByteSize(ts[i]));
             offset += RAI_TensorByteSize(ts[i]);
         }
     } else {
-        status = ort->CreateTensorWithDataAsOrtValue(
-            allocator->Info(allocator), t0->tensor.dl_tensor.data, RAI_TensorByteSize(t0),
-            t0->tensor.dl_tensor.shape, t0->tensor.dl_tensor.ndim,
-            RAI_GetOrtDataTypeFromDL(t0->tensor.dl_tensor.dtype), &out);
-
-        if (status != NULL) {
-            goto error;
-        }
+        ONNX_VALIDATE_STATUS(ort->CreateTensorWithDataAsOrtValue(
+            global_allocator->Info(global_allocator), t0->tensor.dl_tensor.data,
+            RAI_TensorByteSize(t0), t0->tensor.dl_tensor.shape, t0->tensor.dl_tensor.ndim,
+            RAI_GetOrtDataTypeFromDL(t0->tensor.dl_tensor.dtype), &out))
     }
-
     return out;
 
 error:
-    RAI_SetError(error, RAI_EMODELCREATE, ort->GetErrorMessage(status));
+    RAI_SetError(error, RAI_EMODELRUN, ort->GetErrorMessage(status));
     ort->ReleaseStatus(status);
     return NULL;
 }
@@ -154,16 +230,12 @@ RAI_Tensor *RAI_TensorCreateFromOrtValue(OrtValue *v, size_t batch_offset, long
                                          RAI_Error *error) {
     OrtStatus *status = NULL;
     const OrtApi *ort = OrtGetApiBase()->GetApi(1);
-
     RAI_Tensor *ret = NULL;
     int64_t *shape = NULL;
     int64_t *strides = NULL;
 
     int is_tensor;
-    status = ort->IsTensor(v, &is_tensor);
-    if (status != NULL)
-        goto error;
-
+    ONNX_VALIDATE_STATUS(ort->IsTensor(v, &is_tensor))
     if (!is_tensor) {
         // TODO: if not tensor, flatten the data structure (sequence or map) and store it in a
         // tensor. If return value is string, emit warning.
@@ -171,30 +243,18 @@ RAI_Tensor *RAI_TensorCreateFromOrtValue(OrtValue *v, size_t batch_offset, long
     }
 
     ret = RAI_TensorNew();
-
     DLDevice device = (DLDevice){.device_type = kDLCPU, .device_id = 0};
 
     OrtTensorTypeAndShapeInfo *info;
-    status = ort->GetTensorTypeAndShape(v, &info);
-    if (status != NULL)
-        goto error;
+    ONNX_VALIDATE_STATUS(ort->GetTensorTypeAndShape(v, &info))
 
     {
         size_t ndims;
-        status = ort->GetDimensionsCount(info, &ndims);
-        if (status != NULL)
-            goto error;
-
+        ONNX_VALIDATE_STATUS(ort->GetDimensionsCount(info, &ndims))
         int64_t dims[ndims];
-        status = ort->GetDimensions(info, dims, ndims);
-        if (status != NULL)
-            goto error;
-
+        ONNX_VALIDATE_STATUS(ort->GetDimensions(info, dims, ndims))
         enum ONNXTensorElementDataType ort_dtype;
-        status = ort->GetTensorElementType(info, &ort_dtype);
-        if (status != NULL)
-            goto error;
-
+        ONNX_VALIDATE_STATUS(ort->GetTensorElementType(info, &ort_dtype))
         int64_t total_batch_size = dims[0];
         total_batch_size = total_batch_size > 0 ? total_batch_size : 1;
 
@@ -213,23 +273,14 @@ RAI_Tensor *RAI_TensorCreateFromOrtValue(OrtValue *v, size_t batch_offset, long
             strides[i] *= strides[i + 1] * shape[i + 1];
         }
 
-        // size_t sample_bytesize = TF_TensorByteSize(tensor) / total_batch_size;
-
         DLDataType dtype = RAI_GetDLDataTypeFromORT(ort_dtype);
 #ifdef RAI_COPY_RUN_OUTPUT
         char *ort_data;
-        status = ort->GetTensorMutableData(v, (void **)&ort_data);
-        if (status != NULL) {
-            goto error;
-        }
+        ONNX_VALIDATE_STATUS(ort->GetTensorMutableData(v, (void **)&ort_data))
         size_t elem_count;
-        status = ort->GetTensorShapeElementCount(info, &elem_count);
-        if (status != NULL) {
-            goto error;
-        }
+        ONNX_VALIDATE_STATUS(ort->GetTensorShapeElementCount(info, &elem_count))
 
         const size_t len = dtype.bits * elem_count / 8;
-
         const size_t total_bytesize = len * sizeof(char);
         const size_t sample_bytesize = total_bytesize / total_batch_size;
         const size_t batch_bytesize = sample_bytesize * batch_size;
@@ -275,114 +326,70 @@ RAI_Tensor *RAI_TensorCreateFromOrtValue(OrtValue *v, size_t batch_offset, long
     return NULL;
 }
 
-OrtEnv *env = NULL;
-
 RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_ModelOpts opts,
                               const char *modeldef, size_t modellen, RAI_Error *error) {
-    const OrtApi *ort = OrtGetApiBase()->GetApi(1);
 
-    RAI_Device device;
-    int64_t deviceid;
+    const OrtApi *ort = OrtGetApiBase()->GetApi(1);
     char **inputs_ = NULL;
     char **outputs_ = NULL;
-
-    if (!parseDeviceStr(devicestr, &device, &deviceid)) {
-        RAI_SetError(error, RAI_EMODELCREATE, "ERR unsupported device");
-        return NULL;
-    }
-
-    if (deviceid == -1) {
-        // ORT does not like device id as -1
-        deviceid = 0;
-    }
-
+    OrtSessionOptions *session_options = NULL;
+    OrtSession *session = NULL;
     OrtStatus *status = NULL;
 
+    // In the first time we set a model for onnx, we create an environment and register
+    // an allocator to it that uses Redis allocator. This allocator is going to be used for
+    // allocating buffers when creating and running models that run on CPU, and for allocations of
+    // models inputs and outputs names (for both models that run on CPU and GPU)
     if (env == NULL) {
-        status = ort->CreateEnv(ORT_LOGGING_LEVEL_WARNING, "test", &env);
+        ONNX_VALIDATE_STATUS(ort->CreateEnv(ORT_LOGGING_LEVEL_WARNING, "test", &env))
+        ONNX_VALIDATE_STATUS(ort->CreateCustomDeviceAllocator(
+            ORT_API_VERSION, AllocatorAlloc, AllocatorFree, AllocatorInfo, &global_allocator))
+        ONNX_VALIDATE_STATUS(ort->RegisterCustomDeviceAllocator(env, global_allocator))
     }
 
-    if (status != NULL || env == NULL) {
-        goto error;
+    ONNX_VALIDATE_STATUS(ort->CreateSessionOptions(&session_options))
+    if (strcasecmp(devicestr, "CPU") == 0) {
+        // These are required to ensure that onnx will use the registered REDIS allocator (for
+        // a model that defined to run on CPU).
+        ONNX_VALIDATE_STATUS(
+            ort->AddSessionConfigEntry(session_options, "session.use_env_allocators", "1"))
+        ONNX_VALIDATE_STATUS(ort->DisableCpuMemArena(session_options))
     }
 
     // TODO: these options could be configured at the AI.CONFIG level
-    OrtSessionOptions *session_options;
-    status = ort->CreateSessionOptions(&session_options);
-    if (status != NULL) {
-        goto error;
-    }
-
-    // status = ort->SetSessionThreadPoolSize(session_options, 1);
-    if (status != NULL) {
-        ort->ReleaseSessionOptions(session_options);
-        goto error;
-    }
-
-    status = ort->SetSessionGraphOptimizationLevel(session_options, 1);
-    if (status != NULL) {
+    ONNX_VALIDATE_STATUS(ort->SetSessionGraphOptimizationLevel(session_options, ORT_ENABLE_BASIC))
+    ONNX_VALIDATE_STATUS(
+        ort->SetIntraOpNumThreads(session_options, (int)opts.backends_intra_op_parallelism))
+    ONNX_VALIDATE_STATUS(
+        ort->SetInterOpNumThreads(session_options, (int)opts.backends_inter_op_parallelism))
+
+    // If the model is set for GPU, this will set CUDA provider for the session,
+    // so that onnx will use its own allocator for CUDA (not Redis allocator)
+    if (!setDeviceId(devicestr, session_options, error)) {
         ort->ReleaseSessionOptions(session_options);
-        goto error;
-    }
-    ort->SetIntraOpNumThreads(session_options, (int)opts.backends_intra_op_parallelism);
-    ort->SetInterOpNumThreads(session_options, (int)opts.backends_inter_op_parallelism);
-
-    // TODO: we will need to propose a more dynamic way to request a specific provider,
-    // e.g. given the name, in ONNXRuntime
-#if RAI_ONNXRUNTIME_USE_CUDA
-    if (device == RAI_DEVICE_GPU) {
-        OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, deviceid);
-    }
-#else
-    // TODO: Do dynamic device/provider check with GetExecutionProviderType or something on these
-    // lines
-    if (device == RAI_DEVICE_GPU) {
-        RAI_SetError(error, RAI_EMODELCREATE, "ERR GPU requested but ONNX couldn't find CUDA");
         return NULL;
     }
-#endif
-
-    OrtSession *session;
-    status = ort->CreateSessionFromArray(env, modeldef, modellen, session_options, &session);
 
+    ONNX_VALIDATE_STATUS(
+        ort->CreateSessionFromArray(env, modeldef, modellen, session_options, &session))
     ort->ReleaseSessionOptions(session_options);
 
-    if (status != NULL) {
-        goto error;
-    }
-
     size_t n_input_nodes;
-    status = ort->SessionGetInputCount(session, &n_input_nodes);
-    if (status != NULL) {
-        goto error;
-    }
-
+    ONNX_VALIDATE_STATUS(ort->SessionGetInputCount(session, &n_input_nodes))
     size_t n_output_nodes;
-    status = ort->SessionGetOutputCount(session, &n_output_nodes);
-    if (status != NULL) {
-        goto error;
-    }
-
-    OrtAllocator *allocator;
-    status = ort->GetAllocatorWithDefaultOptions(&allocator);
+    ONNX_VALIDATE_STATUS(ort->SessionGetOutputCount(session, &n_output_nodes))
 
     inputs_ = array_new(char *, n_input_nodes);
     for (long long i = 0; i < n_input_nodes; i++) {
         char *input_name;
-        status = ort->SessionGetInputName(session, i, allocator, &input_name);
-        if (status != NULL) {
-            goto error;
-        }
+        ONNX_VALIDATE_STATUS(ort->SessionGetInputName(session, i, global_allocator, &input_name))
         inputs_ = array_append(inputs_, input_name);
     }
 
     outputs_ = array_new(char *, n_output_nodes);
     for (long long i = 0; i < n_output_nodes; i++) {
         char *output_name;
-        status = ort->SessionGetOutputName(session, i, allocator, &output_name);
-        if (status != NULL) {
-            goto error;
-        }
+        ONNX_VALIDATE_STATUS(ort->SessionGetOutputName(session, i, global_allocator, &output_name))
         outputs_ = array_append(outputs_, output_name);
     }
 
@@ -410,62 +417,69 @@ RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_Mo
 
 error:
     RAI_SetError(error, RAI_EMODELCREATE, ort->GetErrorMessage(status));
+    if (session_options) {
+        ort->ReleaseSessionOptions(session_options);
+    }
     if (inputs_) {
         n_input_nodes = array_len(inputs_);
         for (uint32_t i = 0; i < n_input_nodes; i++) {
-            status = ort->AllocatorFree(allocator, inputs_[i]);
+            status = ort->AllocatorFree(global_allocator, inputs_[i]);
         }
         array_free(inputs_);
     }
     if (outputs_) {
         n_output_nodes = array_len(outputs_);
         for (uint32_t i = 0; i < n_output_nodes; i++) {
-            status = ort->AllocatorFree(allocator, outputs_[i]);
+            status = ort->AllocatorFree(global_allocator, outputs_[i]);
         }
         array_free(outputs_);
     }
+    if (session) {
+        ort->ReleaseSession(session);
+    }
     ort->ReleaseStatus(status);
     return NULL;
 }
 
 void RAI_ModelFreeORT(RAI_Model *model, RAI_Error *error) {
     const OrtApi *ort = OrtGetApiBase()->GetApi(1);
-
-    RedisModule_Free(model->data);
-    RedisModule_Free(model->devicestr);
-    OrtAllocator *allocator;
     OrtStatus *status = NULL;
-    status = ort->GetAllocatorWithDefaultOptions(&allocator);
+
     for (uint32_t i = 0; i < model->ninputs; i++) {
-        status = ort->AllocatorFree(allocator, model->inputs[i]);
+        ONNX_VALIDATE_STATUS(ort->AllocatorFree(global_allocator, model->inputs[i]))
     }
     array_free(model->inputs);
 
     for (uint32_t i = 0; i < model->noutputs; i++) {
-        status = ort->AllocatorFree(allocator, model->outputs[i]);
+        ONNX_VALIDATE_STATUS(ort->AllocatorFree(global_allocator, model->outputs[i]))
     }
     array_free(model->outputs);
 
+    RedisModule_Free(model->devicestr);
+    RedisModule_Free(model->data);
     ort->ReleaseSession(model->session);
-
     model->model = NULL;
     model->session = NULL;
+    return;
+
+error:
+    RAI_SetError(error, RAI_EMODELFREE, ort->GetErrorMessage(status));
+    ort->ReleaseStatus(status);
 }
 
 int RAI_ModelRunORT(RAI_ModelRunCtx **mctxs, RAI_Error *error) {
     const OrtApi *ort = OrtGetApiBase()->GetApi(1);
 
     OrtSession *session = mctxs[0]->model->session;
-
     if (session == NULL) {
         RAI_SetError(error, RAI_EMODELRUN, "ERR ONNXRuntime session was not allocated");
-        return 1;
+        return REDISMODULE_ERR;
     }
 
     const size_t nbatches = array_len(mctxs);
     if (nbatches == 0) {
         RAI_SetError(error, RAI_EMODELRUN, "ERR No batches to run");
-        return 1;
+        return REDISMODULE_ERR;
     }
 
     size_t batch_sizes[nbatches];
@@ -483,26 +497,11 @@ int RAI_ModelRunORT(RAI_ModelRunCtx **mctxs, RAI_Error *error) {
     }
 
     OrtStatus *status = NULL;
-
-    OrtAllocator *allocator;
-    status = ort->GetAllocatorWithDefaultOptions(&allocator);
-    if (status != NULL) {
-        goto error;
-    }
-
     size_t n_input_nodes;
-    status = ort->SessionGetInputCount(session, &n_input_nodes);
-    if (status != NULL) {
-        goto error;
-    }
+    ONNX_VALIDATE_STATUS(ort->SessionGetInputCount(session, &n_input_nodes))
 
     size_t n_output_nodes;
-    status = ort->SessionGetOutputCount(session, &n_output_nodes);
-    if (status != NULL) {
-        goto error;
-    }
-
-    {
+    ONNX_VALIDATE_STATUS(ort->SessionGetOutputCount(session, &n_output_nodes)) {
         const char *input_names[n_input_nodes];
         const char *output_names[n_output_nodes];
 
@@ -516,23 +515,20 @@ int RAI_ModelRunORT(RAI_ModelRunCtx **mctxs, RAI_Error *error) {
             char msg[70];
             sprintf(msg, "ERR Expected %li inputs but got %li", n_input_nodes, ninputs);
             RAI_SetError(error, RAI_EMODELRUN, msg);
-            return 1;
+            return REDISMODULE_ERR;
         }
 
         if (noutputs != n_output_nodes) {
             char msg[70];
             sprintf(msg, "ERR Expected %li outputs but got %li", n_output_nodes, noutputs);
             RAI_SetError(error, RAI_EMODELRUN, msg);
-            return 1;
+            return REDISMODULE_ERR;
         }
 
         for (size_t i = 0; i < n_input_nodes; i++) {
             char *input_name;
-            status = ort->SessionGetInputName(session, i, allocator, &input_name);
-            if (status != NULL) {
-                goto error;
-            }
-
+            ONNX_VALIDATE_STATUS(
+                ort->SessionGetInputName(session, i, global_allocator, &input_name))
             input_names[i] = input_name;
 
             RAI_Tensor *batched_input_tensors[nbatches];
@@ -543,75 +539,43 @@ int RAI_ModelRunORT(RAI_ModelRunCtx **mctxs, RAI_Error *error) {
             inputs[i] = RAI_OrtValueFromTensors(batched_input_tensors, nbatches, error);
             if (error->code != RAI_OK) {
                 ort->ReleaseStatus(status);
-                return 1;
+                return REDISMODULE_ERR;
             }
-
-// TODO: use this to check input dim, shapes
-#if 0
-    OrtTypeInfo* typeinfo;
-    status = OrtSessionGetInputTypeInfo(session, i, &typeinfo);
-    const OrtTensorTypeAndShapeInfo* tensor_info = OrtCastTypeInfoToTensorInfo(typeinfo);
-    ONNXTensorElementDataType type = OrtGetTensorElementType(tensor_info);
-    // printf("Input %zu : type=%d", i, type);
-
-    size_t num_dims = OrtGetDimensionsCount(tensor_info);
-    // printf("Input %zu : num_dims=%zu", i, num_dims);
-    input_node_dims.resize(num_dims);
-    OrtGetDimensions(tensor_info, (int64_t*)input_node_dims.data(), num_dims);
-    for (size_t j = 0; j < num_dims; j++) {
-      // printf("Input %zu : dim %zu=%jd", i, j, input_node_dims[j]);
-    }
-
-    OrtReleaseTypeInfo(typeinfo);
-#endif
         }
 
         for (size_t i = 0; i < n_output_nodes; i++) {
             char *output_name;
-            status = ort->SessionGetOutputName(session, i, allocator, &output_name);
-            if (status != NULL) {
-                goto error;
-            }
-
+            ONNX_VALIDATE_STATUS(
+                ort->SessionGetOutputName(session, i, global_allocator, &output_name))
             output_names[i] = output_name;
             outputs[i] = NULL;
         }
 
-        // ORT_API_STATUS(OrtRun, _Inout_ OrtSession* sess,
-        //                _In_ OrtRunOptions* run_options,
-        //                _In_ const char* const* input_names, _In_ const OrtValue* const* input,
-        //                size_t input_len, _In_ const char* const* output_names, size_t
-        //                output_names_len, _Out_ OrtValue** output);
         OrtRunOptions *run_options = NULL;
-        status = ort->Run(session, run_options, input_names, (const OrtValue *const *)inputs,
-                          n_input_nodes, output_names, n_output_nodes, outputs);
+        ONNX_VALIDATE_STATUS(ort->Run(session, run_options, input_names,
+                                      (const OrtValue *const *)inputs, n_input_nodes, output_names,
+                                      n_output_nodes, outputs));
 
-        if (status) {
-            goto error;
+        for (uint32_t i = 0; i < ninputs; i++) {
+            status = ort->AllocatorFree(global_allocator, (void *)input_names[i]);
+        }
+        for (uint32_t i = 0; i < noutputs; i++) {
+            status = ort->AllocatorFree(global_allocator, (void *)output_names[i]);
         }
 
         for (size_t i = 0; i < n_output_nodes; i++) {
             if (nbatches > 1) {
                 OrtTensorTypeAndShapeInfo *info;
-                status = ort->GetTensorTypeAndShape(outputs[i], &info);
-                if (status != NULL)
-                    goto error;
-
+                ONNX_VALIDATE_STATUS(ort->GetTensorTypeAndShape(outputs[i], &info))
                 size_t ndims;
-                status = ort->GetDimensionsCount(info, &ndims);
-                if (status != NULL)
-                    goto error;
-
+                ONNX_VALIDATE_STATUS(ort->GetDimensionsCount(info, &ndims))
                 int64_t dims[ndims];
-                status = ort->GetDimensions(info, dims, ndims);
-                if (status != NULL)
-                    goto error;
-
+                ONNX_VALIDATE_STATUS(ort->GetDimensions(info, dims, ndims))
                 if (dims[0] != total_batch_size) {
                     RAI_SetError(error, RAI_EMODELRUN,
                                  "ERR Model did not generate the expected batch size");
                     ort->ReleaseStatus(status);
-                    return 1;
+                    return REDISMODULE_ERR;
                 }
 
                 for (size_t b = 0; b < nbatches; b++) {
@@ -619,7 +583,7 @@ int RAI_ModelRunORT(RAI_ModelRunCtx **mctxs, RAI_Error *error) {
                         outputs[i], batch_offsets[b], batch_sizes[b], error);
                     if (error->code != RAI_OK) {
                         ort->ReleaseStatus(status);
-                        return 1;
+                        return REDISMODULE_ERR;
                     }
                     if (output_tensor) {
                         mctxs[b]->outputs[i].tensor = RAI_TensorGetShallowCopy(output_tensor);
@@ -633,7 +597,7 @@ int RAI_ModelRunORT(RAI_ModelRunCtx **mctxs, RAI_Error *error) {
                 RAI_Tensor *output_tensor = RAI_TensorCreateFromOrtValue(outputs[i], 0, -1, error);
                 if (error->code != RAI_OK) {
                     ort->ReleaseStatus(status);
-                    return 1;
+                    return REDISMODULE_ERR;
                 }
                 if (output_tensor) {
                     mctxs[0]->outputs[i].tensor = RAI_TensorGetShallowCopy(output_tensor);
@@ -643,21 +607,18 @@ int RAI_ModelRunORT(RAI_ModelRunCtx **mctxs, RAI_Error *error) {
                            "unsupported)");
                 }
             }
-
             ort->ReleaseValue(outputs[i]);
         }
-
         for (size_t i = 0; i < n_input_nodes; i++) {
             ort->ReleaseValue(inputs[i]);
         }
-
-        return 0;
+        return REDISMODULE_OK;
     }
 
 error:
     RAI_SetError(error, RAI_EMODELRUN, ort->GetErrorMessage(status));
     ort->ReleaseStatus(status);
-    return 1;
+    return REDISMODULE_ERR;
 }
 
 int RAI_ModelSerializeORT(RAI_Model *model, char **buffer, size_t *len, RAI_Error *error) {
@@ -665,7 +626,7 @@ int RAI_ModelSerializeORT(RAI_Model *model, char **buffer, size_t *len, RAI_Erro
     memcpy(*buffer, model->data, model->datalen);
     *len = model->datalen;
 
-    return 0;
+    return REDISMODULE_OK;
 }
 
 const char *RAI_GetBackendVersionORT(void) { return OrtGetApiBase()->GetVersionString(); }
diff --git a/src/backends/onnxruntime.h b/src/backends/onnxruntime.h
index ceb233a49..1f68c0051 100644
--- a/src/backends/onnxruntime.h
+++ b/src/backends/onnxruntime.h
@@ -6,6 +6,10 @@
 #include "model_struct.h"
 #include "err.h"
 
+unsigned long long RAI_GetMemoryInfoORT(void);
+
+unsigned long long RAI_GetMemoryAccessORT(void);
+
 int RAI_InitBackendORT(int (*get_api_fn)(const char *, void *));
 
 RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_ModelOpts opts,
diff --git a/src/redisai.c b/src/redisai.c
index 28d016564..af1403e95 100644
--- a/src/redisai.c
+++ b/src/redisai.c
@@ -1102,8 +1102,18 @@ void RAI_moduleInfoFunc(RedisModuleInfoCtx *ctx, int for_crash_report) {
     RedisModule_InfoAddFieldLongLong(ctx, "threads_per_queue", perqueueThreadPoolSize);
     RedisModule_InfoAddFieldLongLong(ctx, "inter_op_parallelism", getBackendsInterOpParallelism());
     RedisModule_InfoAddFieldLongLong(ctx, "intra_op_parallelism", getBackendsIntraOpParallelism());
-    struct rusage self_ru, c_ru;
+    RedisModule_InfoAddSection(ctx, "memory_usage");
+    if (RAI_backends.onnx.get_memory_info) {
+        RedisModule_InfoAddFieldULongLong(ctx, "onnxruntime_memory",
+                                          RAI_backends.onnx.get_memory_info());
+        RedisModule_InfoAddFieldULongLong(ctx, "onnxruntime_memory_access_num",
+                                          RAI_backends.onnx.get_memory_access_num());
+    } else {
+        RedisModule_InfoAddFieldULongLong(ctx, "onnxruntime_memory", 0);
+        RedisModule_InfoAddFieldULongLong(ctx, "onnxruntime_memory_access_num", 0);
+    }
 
+    struct rusage self_ru, c_ru;
     // Return resource usage statistics for the calling process,
     // which is the sum of resources used by all threads in the
     // process
diff --git a/tests/flow/test_data/mul_1.onnx b/tests/flow/test_data/mul_1.onnx
new file mode 100644
index 000000000..c34899b6b
--- /dev/null
+++ b/tests/flow/test_data/mul_1.onnx
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:71f431c4e9321ec6fbeb158d02ed240459a7dcc98673fa79a4f439ce42efaf10
+size 130
diff --git a/tests/flow/tests_onnx.py b/tests/flow/tests_onnx.py
index 849920560..32b5b1fc7 100644
--- a/tests/flow/tests_onnx.py
+++ b/tests/flow/tests_onnx.py
@@ -16,7 +16,6 @@ def test_onnx_modelrun_mnist(env):
         return
 
     con = env.getConnection()
-
     test_data_path = os.path.join(os.path.dirname(__file__), 'test_data')
     model_filename = os.path.join(test_data_path, 'mnist.onnx')
     wrong_model_filename = os.path.join(test_data_path, 'graph.pb')
@@ -452,3 +451,138 @@ def test_parallelism():
                         for k in con.execute_command("INFO MODULES").decode().split("#")[3].split()[1:]}
     env.assertEqual(load_time_config["ai_inter_op_parallelism"], "2")
     env.assertEqual(load_time_config["ai_intra_op_parallelism"], "2")
+
+
+def test_onnx_use_custom_allocator(env):
+    if not TEST_ONNX:
+        env.debugPrint("skipping {} since TEST_ONNX=0".format(sys._getframe().f_code.co_name), force=True)
+        return
+
+    con = env.getConnection()
+    test_data_path = os.path.join(os.path.dirname(__file__), 'test_data')
+    model_filename = os.path.join(test_data_path, 'mul_1.onnx')
+    with open(model_filename, 'rb') as f:
+        model_pb = f.read()
+    ai_memory_config = {k.split(":")[0]: k.split(":")[1]
+                            for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory"]), 0)
+
+    # Expect using the allocator during model set for allocating the model, its input name and output name:
+    # overall 3 allocations. The model raw size is 130B ,and the names are 2B each. In practice we allocate
+    # more than 134B as Redis allocator will use additional memory for its internal management and for the
+    # 64-Byte alignment. When the test runs with valgrind, redis will use malloc for the allocations
+    # (hence will not use additional memory).
+    ret = con.execute_command('AI.MODELSET', 'm{1}', 'ONNX', 'CPU', 'BLOB', model_pb)
+    env.assertEqual(ret, b'OK')
+    ai_memory_config = {k.split(":")[0]: k.split(":")[1]
+                        for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
+
+    # Expect using at least 130+63+(size of an address) + 2*(2+63+(size of an address)) bytes.
+    model_allocation_bytes_used = int(ai_memory_config["ai_onnxruntime_memory"])
+    env.assertTrue(model_allocation_bytes_used > 334)
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory_access_num"]), 3)
+    con.execute_command('AI.TENSORSET', 'a_mul{1}', 'FLOAT', 3, 2, 'VALUES', 1.0, 2.0, 3.0, 4.0, 5.0, 6.0)
+
+    # Running the model should access the allocator 6 times: allocating+freeing input+output names,
+    # and allocating+freeing the output as OrtValue.
+    con.execute_command('AI.MODELRUN', 'm{1}', 'INPUTS', 'a_mul{1}', 'OUTPUTS', 'b{1}')
+    values = con.execute_command('AI.TENSORGET', 'b{1}', 'VALUES')
+    env.assertEqual(values, [b'1', b'4', b'9', b'16', b'25', b'36'])
+    ai_memory_config = {k.split(":")[0]: k.split(":")[1]
+                        for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory_access_num"]), 9)
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory"]), model_allocation_bytes_used)
+
+    # Expect using the allocator free function 3 times: when releasing the model, input name and output name.
+    con.execute_command('AI.MODELDEL', 'm{1}')
+    env.assertFalse(con.execute_command('EXISTS', 'm{1}'))
+    ai_memory_config = {k.split(":")[0]: k.split(":")[1]
+                        for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory"]), 0)
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory_access_num"]), 12)
+
+    # test the use of Redis allocator in model run op.
+    model_filename = os.path.join(test_data_path, 'mnist.onnx')
+    sample_filename = os.path.join(test_data_path, 'one.raw')
+
+    with open(model_filename, 'rb') as f:
+        model_pb = f.read()
+    with open(sample_filename, 'rb') as f:
+        sample_raw = f.read()
+
+    ret = con.execute_command('AI.MODELSET', 'm{1}', 'ONNX', 'CPU', 'BLOB', model_pb)
+    env.assertEqual(ret, b'OK')
+    con.execute_command('AI.TENSORSET', 'a{1}', 'FLOAT', 1, 1, 28, 28, 'BLOB', sample_raw)
+
+    # Expect 18 allocator's access from onnx during the run (in addition to the allocations that were made while
+    # creating the model).
+    ai_memory_config = {k.split(":")[0]: k.split(":")[1]
+                        for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
+    allocator_access_num_before = ai_memory_config["ai_onnxruntime_memory_access_num"]
+    con.execute_command('AI.MODELRUN', 'm{1}', 'INPUTS', 'a{1}', 'OUTPUTS', 'b{1}')
+    ai_memory_config = {k.split(":")[0]: k.split(":")[1]
+                        for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
+    allocator_access_num_after = ai_memory_config["ai_onnxruntime_memory_access_num"]
+    env.assertEqual(int(allocator_access_num_after) - int(allocator_access_num_before), 18)
+
+    values = con.execute_command('AI.TENSORGET', 'b{1}', 'VALUES')
+    argmax = max(range(len(values)), key=lambda i: values[i])
+    env.assertEqual(argmax, 1)
+
+
+def test_onnx_use_custom_allocator_with_GPU(env):
+    if not TEST_ONNX:
+        env.debugPrint("skipping {} since TEST_ONNX=0".format(sys._getframe().f_code.co_name), force=True)
+        return
+    if DEVICE == 'CPU':
+        env.debugPrint("skipping {} since this test if for GPU only".format(sys._getframe().f_code.co_name), force=True)
+        return
+
+    con = env.getConnection()
+    test_data_path = os.path.join(os.path.dirname(__file__), 'test_data')
+    model_filename = os.path.join(test_data_path, 'mul_1.onnx')
+    with open(model_filename, 'rb') as f:
+        model_pb = f.read()
+    ai_memory_config = {k.split(":")[0]: k.split(":")[1]
+                        for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory"]), 0)
+
+    # Expect using the allocator during model set for allocating the model, its input name and output name:
+    # overall 3 allocations. The model raw size is 130B ,and the names are 2B each. In practice we allocate
+    # more than 134B as Redis allocator will use additional memory for its internal management and for the
+    # 64-Byte alignment. When the test runs with valgrind, redis will use malloc for the allocations.
+    ret = con.execute_command('AI.MODELSET', 'm_gpu{1}', 'ONNX', DEVICE, 'BLOB', model_pb)
+    env.assertEqual(ret, b'OK')
+
+    # but for GPU, expect using the allocator only for allocating input and output names (not the model itself).
+    ret = con.execute_command('AI.MODELSET', 'm_cpu{1}', 'ONNX', 'CPU', 'BLOB', model_pb)
+    env.assertEqual(ret, b'OK')
+    ai_memory_config = {k.split(":")[0]: k.split(":")[1]
+                        for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
+
+    # Expect using at least 130+63+(size of an address) + 4*(2+63+(size of an address)) bytes.
+    model_allocation_bytes_used = int(ai_memory_config["ai_onnxruntime_memory"])
+    env.assertTrue(model_allocation_bytes_used > 472)
+    env.assertTrue(model_allocation_bytes_used < 705)
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory_access_num"]), 5)
+
+    # Make sure that allocator is not used for running and freeing the GPU model, except for
+    # the input and output names allocations (and deallocations).
+    con.execute_command('AI.TENSORSET', 'a{1}', 'FLOAT', 3, 2, 'VALUES', 1.0, 2.0, 3.0, 4.0, 5.0, 6.0)
+    con.execute_command('AI.MODELRUN', 'm_gpu{1}', 'INPUTS', 'a{1}', 'OUTPUTS', 'b{1}')
+    values = con.execute_command('AI.TENSORGET', 'b{1}', 'VALUES')
+    env.assertEqual(values, [b'1', b'4', b'9', b'16', b'25', b'36'])
+    # Expect that memory usage didn't change, and for another 4 accesses to the allocator (input and output names
+    # allocation and free)
+    ai_memory_config = {k.split(":")[0]: k.split(":")[1]
+                        for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory"]), model_allocation_bytes_used)
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory_access_num"]), 9)
+
+    # Expect only 2 more accesses in delete - for deallocating input and output names
+    con.execute_command('AI.MODELDEL', 'm_gpu{1}')
+    env.assertFalse(con.execute_command('EXISTS', 'm_gpu{1}'))
+    ai_memory_config = {k.split(":")[0]: k.split(":")[1]
+                        for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory_access_num"]), 11)
+