@@ -153,7 +153,7 @@ struct DeferredToDevice : public Deferred {
153
153
154
154
bool generate_mlir (::mlir::OpBuilder &builder, const ::mlir::Location &loc,
155
155
jit::DepManager &dm) override {
156
- auto av = dm.getDependent (builder, _a );
156
+ auto av = dm.getDependent (builder, Registry::get (_a) );
157
157
158
158
auto srcType = av.getType ().dyn_cast <::imex::ndarray::NDArrayType>();
159
159
assert (srcType);
@@ -183,12 +183,12 @@ struct DeferredToDevice : public Deferred {
183
183
void *r_allocated, void *r_aligned, intptr_t r_offset,
184
184
const intptr_t *r_sizes, const intptr_t *r_strides,
185
185
uint64_t *lo_allocated, uint64_t *lo_aligned) {
186
- auto t = mk_tnsr (reinterpret_cast <Transceiver *>( this ->team () ),
187
- _dtype , this ->shape (), l_allocated, l_aligned ,
188
- l_offset, l_sizes, l_strides, o_allocated ,
189
- o_aligned, o_offset, o_sizes, o_strides ,
190
- r_allocated, r_aligned, r_offset, r_sizes ,
191
- r_strides, lo_allocated, lo_aligned);
186
+ auto t = mk_tnsr (this -> guid (), this ->dtype (), this -> shape ( ),
187
+ this -> device () , this ->team (), l_allocated,
188
+ l_aligned, l_offset, l_sizes, l_strides,
189
+ o_allocated, o_aligned, o_offset, o_sizes,
190
+ o_strides, r_allocated, r_aligned, r_offset,
191
+ r_sizes, r_strides, lo_allocated, lo_aligned);
192
192
this ->set_value (std::move (t));
193
193
});
194
194
return false ;
0 commit comments