You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@beam.apache.org by "Andy Ye (Jira)" <ji...@apache.org> on 2022/04/07 14:25:00 UTC
[jira] [Commented] (BEAM-13986) Provide GPU support
[ https://issues.apache.org/jira/browse/BEAM-13986?page=com.atlassian.jira.plugin.system.issuetabpanels:comment-tabpanel&focusedCommentId=17518920#comment-17518920 ]
Andy Ye commented on BEAM-13986:
--------------------------------
We need to test to make sure GPU trained models loaded into both CPU/GPU containers work. Here's some WIP code to test creating cuda state dicts. I need to look more into this.
state_dict = OrderedDict([
('linear.weight', torch.Tensor([[2.0]], device='cuda')),
('linear.bias', torch.Tensor([0.5], device='cuda'))
])
> Provide GPU support
> -------------------
>
> Key: BEAM-13986
> URL: https://issues.apache.org/jira/browse/BEAM-13986
> Project: Beam
> Issue Type: Sub-task
> Components: sdk-py-core
> Reporter: Andy Ye
> Priority: P2
> Labels: run-inference
>
> Pytorch and Tensorflow have GPU trained models. Need to support this in the RunInference classes, and also make sure they get configured in Dataflow.
--
This message was sent by Atlassian Jira
(v8.20.1#820001)