Skip to content

Commit

Permalink
add casting
Browse files Browse the repository at this point in the history
  • Loading branch information
priyakasimbeg committed Dec 7, 2023
1 parent ad9b843 commit d62ff77
Showing 1 changed file with 2 additions and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -121,6 +121,7 @@ def __init__(self,
super().__init__()

if use_layer_norm:
size = int(size)
norm_layer = nn.LayerNorm([out_chans, size, size], eps=1e-06)
else:
norm_layer = nn.InstanceNorm2d(out_chans)
Expand Down Expand Up @@ -156,6 +157,7 @@ def __init__(self,
):
super().__init__()
if use_layer_norm:
size = int(size)
norm_layer = nn.LayerNorm([out_chans, size, size], eps=1e-06)
else:
norm_layer = nn.InstanceNorm2d(out_chans)
Expand Down

0 comments on commit d62ff77

Please sign in to comment.