Skip to content

Commit

Permalink
format models
Browse files Browse the repository at this point in the history
namsaraeva committed Jun 13, 2024
1 parent 6d00811 commit 5e0d349
Showing 1 changed file with 5 additions and 3 deletions.
8 changes: 5 additions & 3 deletions src/sparcscore/ml/models.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
from collections import OrderedDict

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch
from collections import OrderedDict


class VGGBase(nn.Module):
"""
@@ -88,7 +90,7 @@ def make_layers_MLP(self, cfg_MLP, cfg, regression = False):
in_features = out_features

if regression: # if regression is True, make the final layer a single output
linear = (f"MLP_linear_final", nn.Linear(in_features, 1))
linear = ("MLP_linear_final", nn.Linear(in_features, 1))
layers += [linear]
else:
linear = (f"MLP_linear{i}_classes", nn.Linear(in_features, self.num_classes))

0 comments on commit 5e0d349

Please sign in to comment.