Commit 07e64267 authored by Gustaf Ahdritz's avatar Gustaf Ahdritz
Browse files

Standardize code style

parent de07730f
This diff is collapsed.
...@@ -24,6 +24,7 @@ class PairTransition(nn.Module): ...@@ -24,6 +24,7 @@ class PairTransition(nn.Module):
""" """
Implements Algorithm 15. Implements Algorithm 15.
""" """
def __init__(self, c_z, n, chunk_size=4): def __init__(self, c_z, n, chunk_size=4):
""" """
Args: Args:
...@@ -63,7 +64,7 @@ class PairTransition(nn.Module): ...@@ -63,7 +64,7 @@ class PairTransition(nn.Module):
[*, N_res, N_res, C_z] pair embedding update [*, N_res, N_res, C_z] pair embedding update
""" """
# DISCREPANCY: DeepMind forgets to apply the mask in this module. # DISCREPANCY: DeepMind forgets to apply the mask in this module.
if(mask is None): if mask is None:
mask = z.new_ones(z.shape[:-1]) mask = z.new_ones(z.shape[:-1])
# [*, N_res, N_res, 1] # [*, N_res, N_res, 1]
...@@ -73,7 +74,7 @@ class PairTransition(nn.Module): ...@@ -73,7 +74,7 @@ class PairTransition(nn.Module):
z = self.layer_norm(z) z = self.layer_norm(z)
inp = {"z": z, "mask": mask} inp = {"z": z, "mask": mask}
if(self.chunk_size is not None): if self.chunk_size is not None:
z = chunk_layer( z = chunk_layer(
self._transition, self._transition,
inp, inp,
......
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
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