evolutionary-policy-optimization 0.2.10__tar.gz → 0.2.11__tar.gz

{evolutionary_policy_optimization-0.2.10 → evolutionary_policy_optimization-0.2.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.2.10
+Version: 0.2.11
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization

{evolutionary_policy_optimization-0.2.10 → evolutionary_policy_optimization-0.2.11}/evolutionary_policy_optimization/distributed.py

@@ -61,6 +61,8 @@ def has_only_one_value(t):
     return (t == t[0]).all()
 
 def all_gather_variable_dim(t, dim = 0, sizes = None):
+    device = t.device
+
     if not exists(sizes):
         sizes = gather_sizes(t, dim = dim)
 
@@ -77,7 +79,7 @@ def all_gather_variable_dim(t, dim = 0, sizes = None):
     gathered_tensors = torch.cat(gathered_tensors, dim = dim)
     seq = torch.arange(max_size, device = device)
 
-    mask = einx.less('j i -> (i j)', seq, sizes)
+    mask = einx.less('j, i -> (i j)', seq, sizes)
     seq = torch.arange(mask.shape[-1], device = device)
     indices = seq[mask]
 

{evolutionary_policy_optimization-0.2.10 → evolutionary_policy_optimization-0.2.11}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "evolutionary-policy-optimization"
-version = "0.2.10"
+version = "0.2.11"
 description = "EPO - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{evolutionary_policy_optimization-0.2.10 → evolutionary_policy_optimization-0.2.11}/tests/test_epo.py

@@ -58,7 +58,8 @@ def test_create_agent(
         actor_dim = 256,
         actor_mlp_depth = 2,
         critic_dim = 256,
-        critic_mlp_depth = 4
+        critic_mlp_depth = 4,
+        wrap_with_accelerate = False
     )
 
     state = torch.randn(2, 512)