Commit 016bfd45 authored by Nikita Smetanin's avatar Nikita Smetanin
Browse files

Minor updates & optimizations to support ESMFold

parent 4ac9a223
...@@ -752,6 +752,7 @@ class StructureModule(nn.Module): ...@@ -752,6 +752,7 @@ class StructureModule(nn.Module):
device=device, device=device,
requires_grad=False, requires_grad=False,
), ),
persistent=False,
) )
if not hasattr(self, "group_idx"): if not hasattr(self, "group_idx"):
self.register_buffer( self.register_buffer(
...@@ -761,6 +762,7 @@ class StructureModule(nn.Module): ...@@ -761,6 +762,7 @@ class StructureModule(nn.Module):
device=device, device=device,
requires_grad=False, requires_grad=False,
), ),
persistent=False,
) )
if not hasattr(self, "atom_mask"): if not hasattr(self, "atom_mask"):
self.register_buffer( self.register_buffer(
...@@ -771,6 +773,7 @@ class StructureModule(nn.Module): ...@@ -771,6 +773,7 @@ class StructureModule(nn.Module):
device=device, device=device,
requires_grad=False, requires_grad=False,
), ),
persistent=False,
) )
if not hasattr(self, "lit_positions"): if not hasattr(self, "lit_positions"):
self.register_buffer( self.register_buffer(
...@@ -781,6 +784,7 @@ class StructureModule(nn.Module): ...@@ -781,6 +784,7 @@ class StructureModule(nn.Module):
device=device, device=device,
requires_grad=False, requires_grad=False,
), ),
persistent=False,
) )
def torsion_angles_to_frames(self, r, alpha, f): def torsion_angles_to_frames(self, r, alpha, f):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment