Commit
·
cf18ce2
1
Parent(s):
1b1e134
Upload model
Browse files- modeling_gcn.py +1 -1
modeling_gcn.py
CHANGED
|
@@ -52,7 +52,7 @@ class GCNNet(torch.nn.Module):
|
|
| 52 |
x = scatter(x, batch, dim=-2, reduce='mean')
|
| 53 |
x = self.mlp(x)
|
| 54 |
|
| 55 |
-
return x.
|
| 56 |
|
| 57 |
class GCNModel(PreTrainedModel):
|
| 58 |
config_class = GCNConfig
|
|
|
|
| 52 |
x = scatter(x, batch, dim=-2, reduce='mean')
|
| 53 |
x = self.mlp(x)
|
| 54 |
|
| 55 |
+
return x.squeeze(-1)
|
| 56 |
|
| 57 |
class GCNModel(PreTrainedModel):
|
| 58 |
config_class = GCNConfig
|