Commit d5bf7b61 authored by mova's avatar mova
Browse files

fix linear regression model

parent d06884ee
import torch
from torch import nn
from torch_geometric.nn import global_add_pool
from ..config import conf
......@@ -10,9 +9,9 @@ nfeatures = conf.model.dyn_features + conf.model.static_features
class ModelClass(torch.nn.Module):
def __init__(self):
super(ModelClass, self).__init__()
self.end_lin = nn.Linear(1, 1)
self.end_lin = nn.Linear(2, 1)
def forward(self, batch):
x = global_add_pool(batch.x, batch.batch, size=batch.num_graphs)
x = self.end_lin(x)
return x
X = torch.vstack([batch["ECAL_E"], batch["HCAL_E"]]).float().T
X = self.end_lin(X)
return X
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