Skip to content

recognition/__init__.py from 'cuda' to device parameter #46

@jeffatllatos

Description

@jeffatllatos

Can you change this line only so we can use CPU or GPU.

FROM:
model = Backbone(num_layers=50, drop_ratio=0.6, mode='ir_se').to('cuda').eval()

TO:
model = Backbone(num_layers=50, drop_ratio=0.6, mode='ir_se').to(device).eval()

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions