amend

facebookresearch · Jul 30, 2024 · 6ee7943 · 6ee7943
1 parent b7c8613
commit 6ee7943
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/benchmarl/models/gru.py b/benchmarl/models/gru.py
@@ -71,6 +71,7 @@ def __init__(
         device: DEVICE_TYPING,
         compile: bool,
         centralised: bool,
+        share_params: bool,
     ):
         super().__init__()
         self.input_size = input_size
@@ -80,6 +81,9 @@ def __init__(
         self.compile = compile
         self.centralised = centralised
 
+        if not share_params:
+            raise NotImplementedError
+
         if self.centralised:
             input_size = input_size * self.n_agents
 
@@ -216,6 +220,7 @@ def __init__(
                 self.n_agents,
                 self.device,
                 centralised=self.centralised,
+                share_params=self.share_params,
                 compile=self.compile,
             )
         else: