google-research

1
--- external/gsn/models/discriminator.py	2023-03-09 18:13:17.896724732 +0000
2
+++ external_reference/gsn/models/discriminator.py	2023-03-09 18:07:58.458065115 +0000
3
@@ -1,7 +1,9 @@
4
 import math
5
 import torch
6
 import torch.nn as nn
7
-from .layers import ConvLayer2d, ConvResBlock2d, EqualLinear
8
+from external.gsn.models.layers import ConvLayer2d, ConvResBlock2d, EqualLinear
9
+
10
+from torch_utils import persistence
11
 
12
 
13
 class DiscriminatorHead(nn.Module):
14
@@ -48,6 +50,7 @@
15
         return out
16
 
17
 
18
+@persistence.persistent_class
19
 class ConvDecoder(nn.Module):
20
     def __init__(self, in_channel, out_channel, in_res, out_res):
21
         super().__init__()
22