@@ -575,9 +575,9 @@ def __getitem__(self, idx):
575575 img = torch .from_numpy (self .x [idx ])
576576
577577 target = {}
578- target ["boxes" ] = torch .from_numpy (y [idx ]["boxes" ])
579- target ["labels" ] = torch .from_numpy (y [idx ]["labels" ])
580- target ["scores" ] = torch .from_numpy (y [idx ]["scores" ])
578+ target ["boxes" ] = torch .from_numpy (self . y [idx ]["boxes" ])
579+ target ["labels" ] = torch .from_numpy (self . y [idx ]["labels" ])
580+ target ["scores" ] = torch .from_numpy (self . y [idx ]["scores" ])
581581 mask_i = torch .from_numpy (self .mask [idx ])
582582
583583 return img , target , mask_i
@@ -602,19 +602,33 @@ def __getitem__(self, idx):
602602 if isinstance (target , torch .Tensor ):
603603 target = target .to (self .estimator .device )
604604 else :
605- target ["boxes" ] = target ["boxes" ].to (self .estimator .device )
606- target ["labels" ] = target ["labels" ].to (self .estimator .device )
607- target ["scores" ] = target ["scores" ].to (self .estimator .device )
605+ targets = []
606+ for idx in range (target ["boxes" ].shape [0 ]):
607+ targets .append (
608+ {
609+ "boxes" : target ["boxes" ][idx ].to (self .estimator .device ),
610+ "labels" : target ["labels" ][idx ].to (self .estimator .device ),
611+ "scores" : target ["scores" ][idx ].to (self .estimator .device ),
612+ }
613+ )
614+ target = targets
608615 _ = self ._train_step (images = images , target = target , mask = None )
609616 else :
610617 for images , target , mask_i in data_loader :
611618 images = images .to (self .estimator .device )
612619 if isinstance (target , torch .Tensor ):
613620 target = target .to (self .estimator .device )
614621 else :
615- target ["boxes" ] = target ["boxes" ].to (self .estimator .device )
616- target ["labels" ] = target ["labels" ].to (self .estimator .device )
617- target ["scores" ] = target ["scores" ].to (self .estimator .device )
622+ targets = []
623+ for idx in range (target ["boxes" ].shape [0 ]):
624+ targets .append (
625+ {
626+ "boxes" : target ["boxes" ][idx ].to (self .estimator .device ),
627+ "labels" : target ["labels" ][idx ].to (self .estimator .device ),
628+ "scores" : target ["scores" ][idx ].to (self .estimator .device ),
629+ }
630+ )
631+ target = targets
618632 mask_i = mask_i .to (self .estimator .device )
619633 _ = self ._train_step (images = images , target = target , mask = mask_i )
620634
0 commit comments