modify chunk and sh config

fastfold/distributed/comm.py

@@ -87,7 +87,7 @@ def _chunk_gather(tensor: Tensor, dim=-1, chunks=1) -> Tensor:
         chunk_tensor = tensor.chunk(chunks, dim=1)
 
         for i in range(chunks):
-            _chunk_list = [tensor_list[j*4+i] for j in range(gpc.get_world_size(ParallelMode.TENSOR))]
+            _chunk_list = [tensor_list[j*chunks+i] for j in range(gpc.get_world_size(ParallelMode.TENSOR))]
             _chunk_tensor = chunk_tensor[i]
         
             dist.all_gather(list(_chunk_list),
@@ -108,7 +108,7 @@ def _chunk_gather(tensor: Tensor, dim=-1, chunks=1) -> Tensor:
         chunk_tensor = tensor.chunk(chunks, dim=0)
 
         for i in range(chunks):
-            _chunk_list = [tensor_list[j*4+i] for j in range(gpc.get_world_size(ParallelMode.TENSOR))]
+            _chunk_list = [tensor_list[j*chunks+i] for j in range(gpc.get_world_size(ParallelMode.TENSOR))]
             _chunk_tensor = chunk_tensor[i]
         
             dist.all_gather(list(_chunk_list),

inference.sh

@@ -6,6 +6,8 @@
 
 python inference.py target.fasta data/pdb_mmcif/mmcif_files \
     --output_dir ./ \
+    --gpus 2 \
+    --param_path /data/params/params_model_1.npz  \
     --uniref90_database_path data/uniref90/uniref90.fasta \
     --mgnify_database_path data/mgnify/mgy_clusters_2018_12.fa \
     --pdb70_database_path data/pdb70/pdb70 \