@@ -125,8 +125,8 @@ def valid_step(batch_data, run_info):
125
125
imgs_gpu = imgs_gpu .permute (0 , 3 , 1 , 2 ).contiguous ()
126
126
127
127
# HWC
128
- true_np = torch .squeeze (true_np ).to ( "cuda" ). type (torch .int64 )
129
- true_hv = torch .squeeze (true_hv ).to ( "cuda" ). type (torch .float32 )
128
+ true_np = torch .squeeze (true_np ).type (torch .int64 )
129
+ true_hv = torch .squeeze (true_hv ).type (torch .float32 )
130
130
131
131
true_dict = {
132
132
"np" : true_np ,
@@ -135,7 +135,7 @@ def valid_step(batch_data, run_info):
135
135
136
136
if model .module .nr_types is not None :
137
137
true_tp = batch_data ["tp_map" ]
138
- true_tp = torch .squeeze (true_tp ).to ( "cuda" ). type (torch .int64 )
138
+ true_tp = torch .squeeze (true_tp ).type (torch .int64 )
139
139
true_dict ["tp" ] = true_tp
140
140
141
141
# --------------------------------------------------------------
@@ -155,14 +155,14 @@ def valid_step(batch_data, run_info):
155
155
result_dict = { # protocol for contents exchange within `raw`
156
156
"raw" : {
157
157
"imgs" : imgs .numpy (),
158
- "true_np" : true_dict ["np" ].cpu (). numpy (),
159
- "true_hv" : true_dict ["hv" ].cpu (). numpy (),
158
+ "true_np" : true_dict ["np" ].numpy (),
159
+ "true_hv" : true_dict ["hv" ].numpy (),
160
160
"prob_np" : pred_dict ["np" ].cpu ().numpy (),
161
161
"pred_hv" : pred_dict ["hv" ].cpu ().numpy (),
162
162
}
163
163
}
164
164
if model .module .nr_types is not None :
165
- result_dict ["raw" ]["true_tp" ] = true_dict ["tp" ].cpu (). numpy ()
165
+ result_dict ["raw" ]["true_tp" ] = true_dict ["tp" ].numpy ()
166
166
result_dict ["raw" ]["pred_tp" ] = pred_dict ["tp" ].cpu ().numpy ()
167
167
return result_dict
168
168
0 commit comments