diff --git a/models/rfd3/src/rfd3/inference/input_parsing.py b/models/rfd3/src/rfd3/inference/input_parsing.py
index f8b9021..5826de7 100644
--- a/models/rfd3/src/rfd3/inference/input_parsing.py
+++ b/models/rfd3/src/rfd3/inference/input_parsing.py
@@ -611,7 +611,9 @@ class DesignInputSpecification(BaseModel):
             ).astype(int)
 
             # ... Subset to residues only
-            atom_array_in = atom_array_in[atom_array_in.is_protein]
+            atom_array_in = atom_array_in[
+                atom_array_in.is_protein | atom_array_in.is_dna | atom_array_in.is_rna
+            ]
 
             # ... Set chain ID for unindexed residues as whatever the input has
             start_resid = np.max(atom_array_in.res_id) + 1
diff --git a/models/rfd3na/src/rfd3na/inference/input_parsing.py b/models/rfd3na/src/rfd3na/inference/input_parsing.py
index ad3deb9..aaa91f0 100644
--- a/models/rfd3na/src/rfd3na/inference/input_parsing.py
+++ b/models/rfd3na/src/rfd3na/inference/input_parsing.py
@@ -640,7 +640,9 @@ class DesignInputSpecification(BaseModel):
             ).astype(int)
 
             # ... Subset to residues only
-            atom_array_in = atom_array_in[atom_array_in.is_protein]
+            atom_array_in = atom_array_in[
+                atom_array_in.is_protein | atom_array_in.is_dna | atom_array_in.is_rna
+            ]
 
             # ... Set chain ID for unindexed residues as whatever the input has
             start_resid = np.max(atom_array_in.res_id) + 1