patch

lucidrains · Oct 10, 2023 · 796b242 · 796b242
1 parent 867ab60
commit 796b242
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/setup.py b/setup.py
@@ -4,7 +4,7 @@
   name = 'x-clip',
   packages = find_packages(exclude=[]),
   include_package_data = True,
-  version = '0.14.1',
+  version = '0.14.2',
   license='MIT',
   description = 'X-CLIP',
   author = 'Phil Wang',

diff --git a/x_clip/distributed.py b/x_clip/distributed.py
@@ -4,6 +4,11 @@
 
 from einops import rearrange
 
+def pad_dim_to(t, length, dim = 0):
+    pad_length = length - t.shape[dim]
+    zero_pairs = (-dim - 1) if dim < 0 else (t.ndim - dim - 1)
+    return F.pad(t, (*((0, 0) * zero_pairs), 0, pad_length))
+
 # distributed helpers
 
 def all_gather_variable_dim(t, dim = 0, sizes = None):