Commit 94b631c2 authored by connor.hainje@pnnl.gov's avatar connor.hainje@pnnl.gov
Browse files

Formatting

parent 259667e3
...@@ -134,7 +134,7 @@ class SimpleNet(GeneralNet): ...@@ -134,7 +134,7 @@ class SimpleNet(GeneralNet):
self.n_class = n_class self.n_class = n_class
self.n_detector = n_detector self.n_detector = n_detector
self.fcs = nn.ModuleList( self.fcs = nn.ModuleList(
[nn.Linear(self.n_detector, 1, bias=False) for i in range(self.n_class)] [nn.Linear(self.n_detector, 1, bias=False) for _ in range(self.n_class)]
) )
if const_init is not None: if const_init is not None:
...@@ -163,10 +163,7 @@ class SimpleNet(GeneralNet): ...@@ -163,10 +163,7 @@ class SimpleNet(GeneralNet):
type. type.
""" """
n = self.n_detector n = self.n_detector
outs = [ outs = [self.fcs[i](x[:, i * n : (i + 1) * n]) for i in range(self.n_class)]
self.fcs[i](x[:, i * n : (i + 1) * n])
for i in range(self.n_class)
]
out = torch.cat(outs, dim=1) out = torch.cat(outs, dim=1)
return F.softmax(out, dim=1) return F.softmax(out, dim=1)
......
Supports Markdown
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