REGISTER_OP("Forward")
.Attr("F: {float, int32}")
.Input("img0: F")
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
// c->set_output(0, c->Matrix((3,4,3,3)));
c->set_output(0, c->input(0)); // N C H W
return Status::OK();})
);
How to change the output channels C? I can only set matrix(x,y) of two dimensions, but I need 4 dimensions as the c->input(0) tensor.
Aucun commentaire:
Enregistrer un commentaire