diff --git a/terraform/environments/analytical-platform-compute/kms-keys.tf b/terraform/environments/analytical-platform-compute/kms-keys.tf
index b9bda55fced..7ee087235b5 100644
--- a/terraform/environments/analytical-platform-compute/kms-keys.tf
+++ b/terraform/environments/analytical-platform-compute/kms-keys.tf
@@ -465,19 +465,3 @@ module "actions_runner_cache_efs_kms" {
 
   tags = local.tags
 }
-
-module "sagemaker_kms" {
-  #checkov:skip=CKV_TF_1:Module registry does not support commit hashes for versions
-  #checkov:skip=CKV_TF_2:Module registry does not support tags for versions
-
-  source  = "terraform-aws-modules/kms/aws"
-  version = "3.1.1"
-
-  aliases               = ["sagemaker"]
-  description           = "KMS key for SageMaker encryption-at-rest"
-  enable_default_policy = true
-
-  deletion_window_in_days = 7
-
-  tags = local.tags
-}
diff --git a/terraform/environments/analytical-platform-compute/sagemaker-probation-search.tf b/terraform/environments/analytical-platform-compute/sagemaker-probation-search.tf
index 642b2648a87..7248717edad 100644
--- a/terraform/environments/analytical-platform-compute/sagemaker-probation-search.tf
+++ b/terraform/environments/analytical-platform-compute/sagemaker-probation-search.tf
@@ -6,7 +6,7 @@ locals {
     analytical-platform-compute-development = {
       hmpps-probation-search-dev = {
         namespace       = "hmpps-probation-search-dev"
-        instance_type   = "ml.t2.medium"
+        instance_type   = "ml.m6g.large"
         repository_name = "tei-cpu"
         image_tag       = "2.0.1-tei1.2.3-cpu-py310-ubuntu22.04"
         environment = {
@@ -17,7 +17,7 @@ locals {
     analytical-platform-compute-production = {
       hmpps-probation-search-preprod = {
         namespace       = "hmpps-probation-search-preprod"
-        instance_type   = "ml.t2.medium"
+        instance_type   = "ml.m6g.large"
         repository_name = "tei-cpu"
         image_tag       = "2.0.1-tei1.2.3-cpu-py310-ubuntu22.04"
         environment = {
@@ -59,9 +59,9 @@ resource "aws_sagemaker_model" "probation_search_huggingface_embedding_model" {
 }
 
 resource "aws_sagemaker_endpoint_configuration" "probation_search_config" {
+  #checkov:skip=CKV_AWS_98:KMS key is not supported for NVMe instance storage.
   for_each    = tomap(local.probation_search_environment)
   name        = "${each.value.namespace}-sagemaker-endpoint-config"
-  kms_key_arn = module.sagemaker_kms.key_arn
   production_variants {
     variant_name           = "AllTraffic"
     model_name             = aws_sagemaker_model.probation_search_huggingface_embedding_model[each.key].name