Skip to content

Commit 9ae7b0b

Browse files
Sukhil PatelSukhil Patel
authored andcommitted
Revert batch key change
1 parent df806ee commit 9ae7b0b

File tree

1 file changed

+3
-2
lines changed

1 file changed

+3
-2
lines changed

pvnet/models/multimodal/site_encoders/encoders.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
import einops
66
import torch
7+
from ocf_datapipes.batch import BatchKey
78
from torch import nn
89

910
from pvnet.models.multimodal.linear_networks.networks import ResFCNet2
@@ -74,13 +75,13 @@ def __init__(
7475
)
7576

7677
def _calculate_attention(self, x):
77-
gsp_ids = x["gsp_id"].squeeze().int()
78+
gsp_ids = x[BatchKey.gsp_id].squeeze().int()
7879
attention = self._attention_network(gsp_ids)
7980
return attention
8081

8182
def _encode_value(self, x):
8283
# Shape: [batch size, sequence length, PV site]
83-
pv_site_seqs = x["pv"].float()
84+
pv_site_seqs = x[BatchKey.pv].float()
8485
batch_size = pv_site_seqs.shape[0]
8586

8687
pv_site_seqs = pv_site_seqs.swapaxes(1, 2).flatten(0, 1)

0 commit comments

Comments
 (0)