diff --git a/Cargo.lock b/Cargo.lock
index e935c9c..6b4bf24 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -506,8 +506,8 @@ checksum = "222a222a5bfe1bba4a77b45ec488a741b3cb8872e5e499451fd7d0129c9c7c3d"
 
 [[package]]
 name = "tfhe"
-version = "0.3.0-beta.0"
-source = "git+https://github.com/zama-ai/tfhe-rs.git?rev=0.3.0-beta.0#e810b42eb610b84a544c402fff5fcd4ee789235b"
+version = "0.3.1"
+source = "git+https://github.com/zama-ai/tfhe-rs.git?rev=0.3.1#c58b0a3f6884e2af3110fd24505eb2f8bdb74ddb"
 dependencies = [
  "aligned-vec",
  "bytemuck",
diff --git a/Cargo.toml b/Cargo.toml
index cf0a406..64b559d 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -12,4 +12,4 @@ publish = false
 bincode = "1.3.3"
 clap = { version = "3.1.18", features = ["derive"] }
 serde = "1.0"
-tfhe = { git = "https://github.com/zama-ai/tfhe-rs.git", rev = "0.3.0-beta.0", features = ["boolean", "experimental-force_fft_algo_dif4", "shortint", "integer"] }
+tfhe = { git = "https://github.com/zama-ai/tfhe-rs.git", rev = "0.3.1", features = ["boolean", "experimental-force_fft_algo_dif4", "shortint", "integer"] }
diff --git a/src/gen_keys.rs b/src/gen_keys.rs
index abc258e..2faec22 100644
--- a/src/gen_keys.rs
+++ b/src/gen_keys.rs
@@ -4,13 +4,14 @@
 // All rights reserved.
 
 use tfhe::{
-    generate_keys, shortint::parameters::PARAM_SMALL_MESSAGE_2_CARRY_2_COMPACT_PK, ClientKey,
-    CompactPublicKey, ConfigBuilder, ServerKey,
+    generate_keys,
+    shortint::parameters::parameters_compact_pk::PARAM_MESSAGE_2_CARRY_2_COMPACT_PK_PBS_KS,
+    ClientKey, CompactPublicKey, ConfigBuilder, ServerKey,
 };
 
 pub fn gen_keys() -> (ClientKey, ServerKey, CompactPublicKey) {
     let config = ConfigBuilder::all_disabled()
-        .enable_custom_integers(PARAM_SMALL_MESSAGE_2_CARRY_2_COMPACT_PK, None)
+        .enable_custom_integers(PARAM_MESSAGE_2_CARRY_2_COMPACT_PK_PBS_KS, None)
         .build();
     let (cks, sks) = generate_keys(config);
     let pks = CompactPublicKey::new(&cks);