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):
self.n_class = n_class
self.n_detector = n_detector
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:
......@@ -163,10 +163,7 @@ class SimpleNet(GeneralNet):
type.
"""
n = self.n_detector
outs = [
self.fcs[i](x[:, i * n : (i + 1) * n])
for i in range(self.n_class)
]
outs = [self.fcs[i](x[:, i * n : (i + 1) * n]) for i in range(self.n_class)]
out = torch.cat(outs, 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